diff --git a/.dockerignore b/.dockerignore index 4627a217ba..a54eeb592d 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,7 +1,8 @@ build*/ docs/ test/ +tests/ .cache/ *.swp -models/ \ No newline at end of file +models/ diff --git a/.github/workflows/unsloth-sd-prebuilt.yml b/.github/workflows/unsloth-sd-prebuilt.yml index 4997219464..a8956f791c 100644 --- a/.github/workflows/unsloth-sd-prebuilt.yml +++ b/.github/workflows/unsloth-sd-prebuilt.yml @@ -221,6 +221,11 @@ jobs: # (we build with SD_SERVER_BUILD_FRONTEND / SD_WEBP / SD_WEBM off), so fetch only # ggml to keep the source tarball small and the build self-contained. git submodule update --init --recursive --depth 1 ggml + # The CUDA build includes cmake/cutlass.cmake, which fails without the CUTLASS headers + # (INT8 GEMM and CuTe attention kernels). Only include/ is compiled against, so keep that + # and the license and drop the rest of the checkout from the source tarball. + git submodule update --init --depth 1 thirdparty/cutlass + find thirdparty/cutlass -mindepth 1 -maxdepth 1 ! -name include ! -name LICENSE.txt ! -name .git -exec rm -rf {} + # ggml changes we carry on top of the pinned submodule commit (the submodule points at # leejet/ggml, which we do not push to). A patch that no longer applies fails the build. for p in scripts/unsloth/ggml-patches/*.patch; do @@ -264,7 +269,7 @@ jobs: matrix: include: - label: Darwin-macOS-arm64 - runner: macos-14 + runner: macos-15 arch: arm64 # Static: apple_m1 variant drops FP16_VA (~40% slower). BLAS explicit: sd.cpp routes big matmuls to it. defines: "-DSD_METAL=ON -DGGML_METAL_EMBED_LIBRARY=ON -DGGML_BLAS=ON -DGGML_BLAS_VENDOR=Apple -DCMAKE_OSX_DEPLOYMENT_TARGET=14.0" diff --git a/.gitignore b/.gitignore index b55ea0f81b..1a4646a5e9 100644 --- a/.gitignore +++ b/.gitignore @@ -2,6 +2,7 @@ build*/ cmake-build-*/ test/ tests/ +logs/ .vscode/ .idea/ .cache/ diff --git a/.gitmodules b/.gitmodules index a26210d90a..9e1837e12a 100644 --- a/.gitmodules +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index d7bfda5999..8c5866d494 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -225,6 +225,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 +345,46 @@ 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 (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) diff --git a/README.md b/README.md index b4ab2f5ab2..28af00a697 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,7 @@ API and command-line option may change frequently.*** - [FLUX.2-dev/FLUX.2-klein](./docs/flux2.md) - [FLUX.3 Image Klein / Klein Base](./docs/flux3.md) - [MOSS-Image-Base / MOSS-Image-SFT](./docs/moss_image.md) + - [Fibo 2.0 Dev](./docs/fibo2.md) - [Lens](./docs/lens.md) - [Chroma](./docs/chroma.md) - [Chroma1-Radiance](./docs/chroma_radiance.md) @@ -58,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) @@ -79,6 +81,7 @@ API and command-line option may change frequently.*** - [Boogu Image Edit](./docs/boogu_image.md) - [Mage-Flow-Edit](./docs/mage_flow.md#image-editing) - [LLaDA-Image Edit](./docs/llada_image.md#image-editing) + - [Fibo 2.0 Dev / Edit Turbo](./docs/fibo2.md#editing-and-turbo) - Video Models - [Wan2.1/Wan2.2](./docs/wan.md) - [MiniMax-H3](./docs/minimax_h3.md) diff --git a/assets/fibo2/dev_edit.png b/assets/fibo2/dev_edit.png new file mode 100644 index 0000000000..fd6e06be13 Binary files /dev/null and b/assets/fibo2/dev_edit.png differ diff --git a/assets/fibo2/t2i.png b/assets/fibo2/t2i.png new file mode 100644 index 0000000000..688ffb36ad Binary files /dev/null and b/assets/fibo2/t2i.png differ diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in new file mode 100644 index 0000000000..2dc1e0f6cd --- /dev/null +++ b/cmake/backend_extensions.h.in @@ -0,0 +1,15 @@ +#ifndef __SD_BACKEND_EXTENSIONS_H__ +#define __SD_BACKEND_EXTENSIONS_H__ + +#include "core/cuda/registration.h" + +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 diff --git a/cmake/cutlass.cmake b/cmake/cutlass.cmake new file mode 100644 index 0000000000..0ec6f32f25 --- /dev/null +++ b/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() diff --git a/cmake/embed_spirv.cmake b/cmake/embed_spirv.cmake new file mode 100644 index 0000000000..a664209693 --- /dev/null +++ b/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 \nnamespace sd_vulkan_shaders {\nstatic const uint32_t ${NAME}[] = {${SPIRV_WORDS}};\n}\n") diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake index d3e952dca4..9370728457 100644 --- a/cmake/ggml.cmake +++ b/cmake/ggml.cmake @@ -9,20 +9,29 @@ 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) -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).") +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() -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}") # The fused DiT ops come from scripts/unsloth/ggml-patches/0004; an unpatched ggml keeps the unfused H3 graph. -set(sd_ggml_header "${sd_ggml_private_include}/../include/ggml.h") +# Upstream only resolves the private include dir for SD_USE_UPSTREAM_GGML now, so find ggml.h here. +get_target_property(sd_ggml_ops_imported ggml IMPORTED) +if(sd_ggml_ops_imported) + set(sd_ggml_header "${SD_GGML_SOURCE_DIR}/include/ggml.h") +else() + get_target_property(sd_ggml_ops_src ggml SOURCE_DIR) + set(sd_ggml_header "${sd_ggml_ops_src}/../include/ggml.h") +endif() if(EXISTS "${sd_ggml_header}") set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS "${sd_ggml_header}") file(STRINGS "${sd_ggml_header}" sd_ggml_h3_fused_ops REGEX "ggml_rope_pe_permute\\(") diff --git a/cmake/hip.cmake b/cmake/hip.cmake new file mode 100644 index 0000000000..e0d813340c --- /dev/null +++ b/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() diff --git a/docker/Dockerfile b/docker/Dockerfile index 3a80ae9bf8..30735094e9 100644 --- a/docker/Dockerfile +++ b/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 && \ diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 625dc60e70..e3364e8402 100644 --- a/docker/Dockerfile.cuda +++ b/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 diff --git a/docker/Dockerfile.musa b/docker/Dockerfile.musa index 08d6f1d43d..9042452d38 100644 --- a/docker/Dockerfile.musa +++ b/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 && \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index fe8a9b15e1..9236bdb1cd 100644 --- a/docker/Dockerfile.rocm +++ b/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/* diff --git a/docker/Dockerfile.sycl b/docker/Dockerfile.sycl index 9225e6eb2f..bdbb7509c4 100644 --- a/docker/Dockerfile.sycl +++ b/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 && \ diff --git a/docker/Dockerfile.vulkan b/docker/Dockerfile.vulkan index c7810a1fa5..576e5d9989 100644 --- a/docker/Dockerfile.vulkan +++ b/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 && \ diff --git a/docs/adetailer.md b/docs/adetailer.md index 2bc64df2c0..f42831823c 100644 --- a/docs/adetailer.md +++ b/docs/adetailer.md @@ -1,9 +1,10 @@ # ADetailer `sd-cli` can run a YOLOv8 object detector on an existing or newly generated -image and perform a cropped inpaint pass for every detected object. The first -implementation supports YOLOv8 detection checkpoints. YOLOv8 segmentation and -MediaPipe models are not supported yet. +image, and `sd-server` can run it after image generation. Both perform a cropped +inpaint pass for every detected object. The implementation supports YOLOv8 +detection checkpoints. YOLOv8 segmentation and MediaPipe models are not +supported yet. ## Convert a detector @@ -68,6 +69,36 @@ ADetailer can also run automatically after normal image generation: --extra-ad-args "confidence=0.3,denoising_strength=0.4,inpaint_width=512,inpaint_height=512" ``` +The same options can be supplied when starting `sd-server`. When `--ad-model` is set, its detector runs after OpenAI-compatible and native asynchronous image generation. SDAPI requests enable ADetailer through `alwayson_scripts`, as described below. If no object is detected, the original image is returned unchanged. RGBA output keeps its original alpha channel, including when detected regions are repaired. Detector initialization or detail-pass failures are reported as request failures. + +Server image-generation requests limit base, hires, and detail-pass dimensions to 4096 pixels per side. ADetailer step overrides are limited to 1–100; `steps=0` in `extra_ad_args` keeps inheriting the main generation steps. These limits also apply to inherited dimensions and `extra_ad_args` overrides. Requests exceeding the limits return a parameter error before generation. CLI limits are unchanged. + +### SDAPI requests + +`/sdapi/v1/txt2img` and `/sdapi/v1/img2img` accept the WebUI ADetailer argument format: + +```json +{ + "prompt": "portrait photo", + "alwayson_scripts": { + "ADetailer": { + "args": [true, false, { + "ad_model": "face_yolov8n.pt", + "ad_prompt": "[PROMPT], detailed face", + "ad_use_steps": true, + "ad_steps": 20 + }] + } + } +} +``` + +The detector must be configured at server startup with `--ad-model`. `ad_model` must match that file's basename; the original `.pt` basename is also accepted for a converted detector. Requests cannot load arbitrary detector paths. Omitting the script, setting its first argument to `false`, selecting `"None"`, or setting `ad_tab_enable` to `false` disables detailing. Dictionary-only argument arrays are also accepted. Multiple enabled dictionaries run in order using the configured detector. + +Supported settings include prompts, confidence, mask filtering by area or confidence, mask count and area ratios, dilation/erosion, offsets, mask merging/inversion, blur, denoising strength, masked or full-image inpainting, padding, and optional overrides for inpaint dimensions, steps, CFG, sampler/scheduler, and CLIP skip. The corresponding `ad_use_*` flags control whether overrides apply. Otherwise dimensions, steps, CFG, sampler, scheduler, and CLIP skip inherit the main request. Denoising strength defaults to `0.4`. Pixel mask settings are limited to the larger output dimension, with a minimum limit of 32 and a maximum of 16384 pixels. + +Model/VAE switching, non-default noise multipliers, face restoration, ControlNet, and skipping the initial img2img pass are not supported; active requests for these features return an error. Unknown settings and invalid active values also return errors. During detailing, `/sdapi/v1/progress` remains active with an `adetailer` job; its percentage and preview stay fixed until detailing finishes. + An empty ADetailer prompt inherits the main prompt. `[PROMPT]` inserts the main prompt, `[SEP]` assigns different prompts to consecutive masks, and `[SKIP]` skips the corresponding mask. @@ -82,6 +113,7 @@ All settings other than the detector path and prompts are passed through | `nms` | `0.45` | NMS IoU threshold | | `max_detections` | `100` | Maximum detections retained after NMS | | `mask_k_largest` | `0` | Keep only the largest K masks; zero keeps all | +| `mask_filter_method` | `area` | Rank retained masks by `area` or `confidence` | | `mask_min_ratio` | `0` | Minimum bbox area relative to the image | | `mask_max_ratio` | `1` | Maximum bbox area relative to the image | | `dilate_erode` | `4` | Positive values dilate; negative values erode | @@ -90,6 +122,7 @@ All settings other than the detector path and prompts are passed through | `merge_masks`, `invert_mask` | `false` | Boolean alternatives to `mask_mode` | | `mask_blur` | `4` | Final composite feather radius | | `inpaint_padding` | `32` | Padding around the detected region | +| `inpaint_only_masked` | `true` | Crop to the detected region; `false` uses the full image with the detection mask | | `inpaint_width`, `inpaint_height` | mode-specific | `512x512` after generation; input/output size in `adetailer` mode | | `denoising_strength` | mode-specific | `0.4` after generation; inherits `--strength` in `adetailer` mode | | `steps` | `0` | Detail steps; zero inherits the main generation | diff --git a/docs/backend.md b/docs/backend.md index 057d3f20ef..019c328171 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -139,7 +139,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. @@ -251,6 +254,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 diff --git a/docs/build.md b/docs/build.md index 6ec0e40f0a..62966ec6f5 100644 --- a/docs/build.md +++ b/docs/build.md @@ -2,18 +2,23 @@ ## 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 ``` ## Selecting a GGML source tree @@ -27,8 +32,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 +52,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 +99,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 +111,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. diff --git a/docs/docker.md b/docs/docker.md index 28e9ab5dfa..e227d859dc 100644 --- a/docs/docker.md +++ b/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" -``` \ No newline at end of file +``` diff --git a/docs/edit.md b/docs/edit.md index 62ecff2af5..1312854b71 100644 --- a/docs/edit.md +++ b/docs/edit.md @@ -19,6 +19,7 @@ Depending on the architecture, different models handle reference images differen | **Krea2 (Community Edit LoRAs)** | `krea2_ostris_edit` | | [**Mage-Flow-Edit**](./mage_flow.md#image-editing) | `mage_flow` | | [**LLaDA-Image**](./llada_image.md#image-editing) | `llada_image` | +| [**Fibo 2.0 Dev / Edit Turbo**](./fibo2.md#examples) | `fibo2` | | **Anima (Community Edit LoRAs)** | `cosmos_reference` | Stable-diffusion.spp also supports basic Unet-based editing models like instruct-pix2pix or CosXL-Edit. This document is not about those. @@ -50,6 +51,7 @@ The `--ref-image-args` argument accepts a comma-separated list of key-value pair | `flux_kontext` | FLUX.1 Kontext | | `longcat` | LongCat Image Edit | | `flux2` | FLUX.2 models | +| `fibo2` | Fibo 2.0 Dev / Edit Turbo | | `qwen` | Qwen Image Edit | | `qwen_layered` | Qwen Image Layered | | `z_image_omni` | Boogu, Z-Image Omni | @@ -73,11 +75,19 @@ If presets are insufficient, you can manually configure the following parameters | `ref_index_mode` | string | Behavior of the RoPE index. | `fixed`, `increase`, `decrease` | | `force_ref_timestep_zero` | bool | Forces timestep=0 for reference tokens. | `true`, `false` (Krea2 only) | | `resize_before_vae` | bool | Whether reference images are resized before VAE encoding. | `true`, `false` | -| `vae_input_max_pixels` | int | Maximum pixel area for VAE reference inputs. | Integer | +| `resize_vae_to_target` | bool | Use output dimensions for every reference, overriding area-based sizing. | `true`, `false` | +| `resize_single_ref_to_target` | bool | Use output dimensions when exactly one reference is supplied. | `true`, `false` | +| `vae_size_buckets` | bool | Choose aligned dimensions within ±10% of the requested area, minimizing aspect-ratio cropping. Ignored when using output dimensions. | `true`, `false` | +| `vae_center_crop` | bool | Scale to cover the selected dimensions with Lanczos, then center-crop. Otherwise resize directly. | `true`, `false` | +| `vae_input_max_pixels` | int | VAE area budget, also capped by output area in ordinary resizing. With size buckets, this is the target area and output area does not cap it. | Integer | +| `vlm_use_vae_input` | bool | Reuse the resized/cropped VAE input for the VLM, before the VLM's own resize. Has no effect when `resize_before_vae=false`. | `true`, `false` | | `vlm_resize_mode` | string | How to resize VLM reference inputs. | `longest_side`, `area`, `none` | | `vlm_max_size` | int | Maximum VLM input size; interpreted according to `vlm_resize_mode`. | Integer | | `vlm_min_size` | int | Minimum VLM input size; interpreted according to `vlm_resize_mode`. | Integer | | `vlm_size` | int | Shortcut to set both VLM min and max size to the same value. | Integer | +| `ref_image_size` | string | MiniMax-H3 reference image sizing before VAE encoding. `match` scales each reference down to the generation's pixel area; `max` uses a 2048 px short edge. Only references above the limit are downscaled; each axis is then snapped to the nearest 32 px multiple (an axis of at least 16 px grows by at most 16 px, and references are never narrower than 32 px). | `match`, `max` | + +`ref_image_size` applies to `vid_gen` with MiniMax-H3 Ref2VA. Reference latents are attended to at every sampling step, so `max` can be several times slower. ### Preset Default Values @@ -88,6 +98,7 @@ For a technical overview of how each preset is configured, see the table below. | `flux_kontext` | No | `fixed` | `none` | | | `longcat` | Yes | `fixed` | `area` | | | `flux2` | No | `increase` | `none` | | +| `fibo2` | Yes | `increase` | `longest` | Shared reference crops; clean reference timesteps; see [Fibo 2.0](./fibo2.md) | | `qwen` | Yes | `increase` | `area` | | | `qwen_layered` | Yes | `decrease` | `area` | | | `mage_flow` | Yes | `increase` | `longest` | `vlm_max_size = 384`, VAE input resized to target | @@ -99,3 +110,5 @@ For a technical overview of how each preset is configured, see the table below. **Additional Default Notes:** - **VLM Input Sizes:** For most presets, `vlm_max_size` and `vlm_min_size` are set to `-1`, meaning the values are model-dependent and handled automatically. In `area` mode they represent pixel area; in `longest_side` mode they represent a side length in pixels. - **VAE Input Size:** `vae_input_max_pixels` defaults to $1024 \times 1024$ pixels (`1048576`). +- **VAE Geometry:** `resize_vae_to_target` defaults to `true` for Mage-Flow and LLaDA-Image. Fibo enables `resize_single_ref_to_target`, `vae_size_buckets`, and `vae_center_crop`. These options are otherwise `false`. Bucket dimensions use the model's alignment (normally 16, Qwen Image 32), with a minimum side of four alignment units; if no bucket fits, output dimensions are used. +- **Shared Encoder Input:** `vlm_use_vae_input` defaults to `true` for Fibo and Qwen Image 2.1, otherwise `false`. All these options can be overridden individually after selecting a preset. diff --git a/docs/fibo2.md b/docs/fibo2.md new file mode 100644 index 0000000000..064d800892 --- /dev/null +++ b/docs/fibo2.md @@ -0,0 +1,53 @@ +# How to Use + +Fibo 2.0 uses a Fibo diffusion transformer, the standard FLUX.2 VAE, and Qwen3-VL 4B as the LLM text encoder. The Qwen tokenizer is built in. + +## Download weights + +- Download Fibo 2.0 Dev + - safetensors: https://huggingface.co/Comfy-Org/Fibo-2.0/tree/main/diffusion_models +- Download Fibo 2.0 Edit Turbo + - safetensors: https://huggingface.co/Comfy-Org/Fibo-2.0/tree/main/diffusion_models +- Download vae + - safetensors: https://huggingface.co/black-forest-labs/FLUX.2-dev/tree/main +- Download Qwen3-VL 4B + - safetensors: https://huggingface.co/Comfy-Org/Fibo-2.0/tree/main/text_encoders + - gguf: https://huggingface.co/Qwen/Qwen3-VL-4B-Instruct-GGUF/tree/main + +The original Diffusers weights from [briaai/Fibo-2.0-dev](https://huggingface.co/briaai/Fibo-2.0-dev) and [briaai/Fibo-2.0-edit-turbo](https://huggingface.co/briaai/Fibo-2.0-edit-turbo) are also supported. Pass the transformer's `diffusion_pytorch_model.safetensors.index.json` to `--diffusion-model` and keep its shards alongside it. + +INT8 convrot weights load directly; see [INT8 convrot](./int8_convrot.md) for backend requirements. + +## Examples + +### Fibo 2.0 Dev + +Dev supports both text-to-image generation and reference-image editing. + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\fibo2_dev_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\flux2_ae.safetensors -p "a lovely cat holding a sign says 'fibo2.cpp'" --cfg-scale 5.0 --sampling-method euler -v --offload-to-cpu +``` + + + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\fibo2_dev_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\flux2_ae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'fibo2.cpp'" --cfg-scale 5.0 --sampling-method euler -v --offload-to-cpu +``` + + + +### Fibo 2.0 Edit Turbo + +Edit Turbo is distilled for reference-image editing in 4 steps with CFG 1. + +``` +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\fibo2_edit_turbo_int8_convrot.safetensors --llm ..\models\text_encoders\Qwen3-VL-4B-Instruct-Q4_K_M.gguf --llm_vision ..\models\text_encoders\Qwen3-VL-4B-Instruct-mmproj-BF16.gguf --vae ..\models\vae\flux2_ae.safetensors -r ..\assets\flux\flux1-dev-q8_0.png -p "change 'flux.cpp' to 'fibo2.cpp'" --cfg-scale 1 --steps 4 --sampling-method euler -v --offload-to-cpu +``` + +The BF16 files accept the same options. Use dimensions divisible by 16. Repeat `-r` to use multiple reference images; instructions can refer to ``, ``, etc. See [Image Editing](./edit.md) for reference-image preprocessing options. Separate GGUF text encoders need their vision projector passed with `--llm_vision` for editing. + +The automatic `fibo2` reference preset combines `resize_single_ref_to_target=true`, `vae_size_buckets=true`, `vae_center_crop=true`, and `vlm_use_vae_input=true`. A single reference is scaled to cover the output dimensions and center-cropped. Multiple references use dimensions divisible by 16, chosen within ±10% of one megapixel to minimize cropping for each image's aspect ratio; `vae_input_max_pixels` overrides that target area. The VAE and vision encoder receive the same crops, with smaller copies used by the vision encoder. These are generic reference-image options and can be overridden individually: for example, `--ref-image-args vae_center_crop=false` uses direct resizing, while `resize_before_vae=false` skips this automatic crop and resize step entirely. + +The vision encoder defaults to `vlm_resize_mode=longest_side,vlm_max_size=768`. Override these with `--ref-image-args`, for example `vlm_max_size=512`, or use `vlm_resize_mode=none` to skip size-based scaling. Vision inputs are still aligned to the encoder's patch grid. + +Plain prompts and structured JSON captions are accepted. This integration uses the caption-only text-rendering path (`render_text_images=False`); it does not generate auxiliary glyph images or expand prompts with Bria VGL. Native Fibo mask preprocessing is not connected to `--mask`. diff --git a/docs/generation_metadata.md b/docs/generation_metadata.md new file mode 100644 index 0000000000..1f7fc3338c --- /dev/null +++ b/docs/generation_metadata.md @@ -0,0 +1,165 @@ +# Generation results and metadata + +`generate_image()`, `generate_video()`, `adetail_image()` and `upscale()` return +owned results containing the media and its generation record. This is a breaking +C API change: applications and bindings must update their calls and rebuild. + +```c +sd_img_gen_params_t params; +sd_img_gen_params_init(¶ms); +params.prompt = "a small house"; + +sd_image_generation_result_t* result = NULL; +if (generate_image(ctx, ¶ms, &result)) { + for (int i = 0; i < result->count; ++i) { + const sd_image_result_t* item = &result->items[i]; + /* Consume item->image and item->generation_info together. */ + } + free_sd_image_generation_result(result); +} +``` + +`result_out` is required and is set to `NULL` on failure. A successful result +owns its pixels and UTF-8 JSON strings. The video result additionally owns its +audio and contains the effective `fps` and `frame_count`. Results remain valid +after freeing the generation context. Use the matching library result destructor; +do not separately free fields unless transferring ownership and clearing them. +Both result destructors accept `NULL`. + +## Reproduction parameters + +The `sdcpp.generation/v1` JSON is a snapshot of generation parameters for the +same loaded context and the same input images, masks, reference media and other +assets. It is created after `GenerationRequest` resolves the request, before +preprocessing or sampling. A failed call returns no result or metadata. + +Fields follow the C generation parameter structs: `prompt`, `negative_prompt`, +`width`, `height`, `seed`, `sample_params`, `hires`, `cache`, +`vae_tiling_params`, `image_preprocess`, and the image/video-specific options. +`kind` distinguishes image and video parameters. There is no inferred `task`. +Enums use their API string names. Seeds are decimal strings to preserve their +full integer range. Omitted optional fields retain the parameter initializer's +defaults, including automatic values represented by infinity in the C API. + +Disabled optional features are omitted entirely: `cache`, `hires`, VAE tiling, +zero/nonpositive `shifted_timestep`, inactive SLG, false circular axes, and empty +preprocessing/reference arguments. SLG layers remain present when `slg_uncond` +is enabled even with zero SLG scale. Enabled caches store their mode's parameters; +TaylorSeer retains the shared DBCache thresholds and limits used by the runtime. +Image CFG at its neutral value of 1 is omitted and resolves back to 1 by default. + +Image strength, ControlNet/IP-Adapter strength and PhotoMaker settings are stored +when their corresponding input images are supplied. PuLID requires an embedding +path. Layer count is stored for layered generation; video VACE and high-noise +settings require those models. High-noise settings omit scheduler, custom sigmas, +flow shift and timestep shifting because that pass shares the base configuration. +Explicit nonempty argument strings are retained without rewriting their contents. +Core sampling values such as seed, sampler, scheduler, steps and guidance remain +explicit even when equal to their defaults. ADetailer always retains its inpaint +strength because its input crops are created later. +Model-required VAE tiling is resolved before serialization, so forced tiling +(such as MiniMax-H3) remains present even if the caller did not enable it. + +The request fixes the seed, dimensions, clip skip, guidance, sampler, scheduler +and eta before taking the snapshot. `custom_sigmas` is stored only when supplied +as an input. Automatically computed schedules, preprocessing transforms, +conditioning/cache decisions, detector regions and output statistics are not +stored: they are recomputed from the same parameters and inputs. + +Each image records the seed used to sample its latent. `batch_count` is omitted; +replay image parameters with `batch_count = 1`. Layers or frames decoded from +the same latent share its seed, so layered generation reproduces that group of +layers. Video frames share the video request. Partial batch cancellation returns +metadata only for completed outputs. Full cancellation returns no result. + +The current image pipeline prepares input latents once using the starting seed +and shares them across a batch. With a stochastic VAE encoder, replaying a later +image alone using its sampling seed can change those input latents. Per-image +seeds therefore do not guarantee independent reproduction for these paths. + +To replay a record, initialize the matching C parameter struct, populate its +fields from the JSON, restore the same input buffers, and call generation with +the same context. Keep custom-sigma/SLG arrays and string storage alive for the +call. Paths supplied as generation parameters (LoRAs, hires models and identity +embeddings) are retained. `models` identifies successfully loaded context model +files by component (`model`, `diffusion_model`, `clip_l`, `vae`, etc.), using +filenames with extensions and omitting directories. Optional components disabled +for the request and preview-only TAE are omitted. WebUI text includes `Model` +(the diffusion model when supplied separately, otherwise the checkpoint, without +the extension) and `VAE` for a separate active VAE. These names are identifiers, +not hashes or sufficient configuration to rebuild a context; replay still requires +the same model files and context settings. The JSON does not embed media buffers +or model weights. + +When `--lora-model-dir` is configured, CLI/server output records express LoRA +paths relative to that directory, preserving subdirectories and file extensions +and using `/` separators. This applies to file metadata, API responses and +ADetailer records. Resolve these paths against the same LoRA directory before +passing them back to the C API. Core results retain the supplied loading paths. +Files outside the directory retain `..` components; paths on another volume, +or records without a configured LoRA directory, retain their original paths. + +## Postprocessing + +Upscale records `upscale.upscale_factor`. ADetailer records `adetailer` and +`inpaint_params`, with a concrete starting seed chosen before detection. These +operations require the same upscaler/detector context and source image on replay. +They do not store detection results or per-region generation traces. + +The CLI/server append operation parameter snapshots to `postprocess` in call +order. The source generation parameters remain unchanged. Frame postprocessing +uses the same operation parameters for each video frame. Standalone +postprocessing does not invent source diffusion parameters for an input image. + +## File metadata + +Image encoders accept the native record and derive WebUI infotext from it. +Images use the existing metadata locations and combined text layout: +`, SDCPP: `. Standalone postprocessing without +source diffusion parameters stores only native JSON. Video metadata stores +only native JSON. + +| Container | Metadata location | +| --- | --- | +| PNG | `tEXt`, keyword `parameters` | +| JPEG | COM segment, `parameters` label followed by a NUL byte | +| WebP image / animated WebP | XMP property `sdcpp:parameters` | +| AVI | RIFF INFO `ICMT` | +| WebM | `COMMENT` tag | + +Inspect images or videos with `sd-cli -M metadata --image `; the same +command supports all containers above, with `--metadata-format json` for a +structured report. Video metadata inspection does not require a model or video +decoder. +Matroska containers are also supported, including `.webm` outputs whose audio +codec causes the muxer to select the Matroska document type. + +The XMP namespace is `https://github.com/leejet/stable-diffusion.cpp/ns/1.0/`. +PNG and JPEG retain the existing UTF-8 text payload convention. Memory encoding +and file saving use the same implementation. The core result and API +`generation_info` fields remain pure JSON; the combined layout applies only +to embedded image metadata. + +WebUI conversion is lossy: known samplers and schedules have explicit name +mappings; unsupported ones use `SDCPP Sampler`/`SDCPP Schedule` fields. Hires +keeps the base request size and separate configured hires size/steps. The native +record retains the full generation parameters. Standalone postprocessing has no WebUI +infotext because its original diffusion settings are unknown. + +WebUI prompt text appends ordinary LoRAs as ``. +Directory components are omitted to match WebUI's filename lookup; the native +`loras` array retains relative paths and weights. High-noise-only LoRAs and names +that cannot be represented by this tag syntax remain in the native record only. +The native prompt is not modified, and postprocessing LoRAs are not added to the +source generation prompt. + +JPEG requires the COM payload, including the label and terminating NUL bytes, +to fit within 65,533 bytes. Larger records fail encoding with an error rather +than being truncated. Use PNG or WebP for larger records. + +Disabling file metadata omits both image representations, but does not disable +the core generation record or its API response. Native async image results and +OpenAI image response items contain `generation_info`; video results contain a +sequence-level `generation_info`. The SD API `info` object contains per-output +`generation_info` alongside derived `infotexts` and each image's seed. Its outer +`parameters` field continues to echo the request. diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md new file mode 100644 index 0000000000..5c059892e8 --- /dev/null +++ b/docs/ggml_extensions.md @@ -0,0 +1,485 @@ +# 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. Intermediate allocations are retained +in this initial implementation; only their computation is removed. + +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 + +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. | + +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. + +## sd.cpp integration + +`src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas, +validation, graph construction helpers, and backend registration entry point. +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. + +`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. + +`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_mul_mat_i8_tensorwise`, and `ggml_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_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_mul_mat_i8_tensorwise` also accepts F32 activations (quantized into temporary +workspace) or the output of `ggml_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. +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. The CUTLASS SwiGLU/ConvRot chain can write its packed +output directly after GEMM, because quantization reads only its workspace +intermediate. 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 existing `ggml_sage_attn` entry point is declared and implemented in sd.cpp's +common extension files. 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 existing `ggml_sol_attn` entry point now lives in the common sd.cpp extension +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. diff --git a/docs/image_preprocessing.md b/docs/image_preprocessing.md index 719781750e..c097e5a901 100644 --- a/docs/image_preprocessing.md +++ b/docs/image_preprocessing.md @@ -100,10 +100,13 @@ for individual references, identity images, or video control frames. --image-preprocess "target=control-frame,index=2,canny=true" ``` -Init and mask sources must have the same dimensions. The mask inherits the -init crop, resize, and padding coordinates, while retaining its own filter, -padding value, and Canny setting. Conflicting mask geometry is rejected. An -omitted mask remains absent until the original pipeline creates its default mask. +If init and mask source dimensions differ, the mask is first center-cropped to +the init source aspect ratio and resized to the init source dimensions using +nearest-exact interpolation. The mask then inherits the init crop, resize, and +padding coordinates, while retaining its own filter, padding value, and Canny +setting for that transform. Explicit mask geometry is evaluated after this +alignment; conflicting geometry is rejected. An omitted mask remains absent +until the original pipeline creates its default mask. ## Downstream behavior @@ -128,7 +131,8 @@ To disable reference resizing before VAE encoding, use `--ref-image-args "resize_before_vae=false"` or the server field `"ref_image_args": "resize_before_vae=false"`. This is separate from `target=ref,mode=none`, which only skips input geometry. Model constraints -still apply. +still apply; MiniMax-H3 requires 32-pixel-aligned references and warns instead +of honoring the switch. ## Server requests @@ -148,12 +152,12 @@ an array of rule strings: OpenAI-compatible requests accept it through `{...}` in the prompt. Request rules replace server-default rules. Generation metadata records the -user rules; image encodings and channel conventions are unchanged. +user rules for replay with the same input images; image encodings and channel conventions are unchanged. ## C API Set `image_preprocess` on the existing image/video generation parameters. -The `generate_image()` and `generate_video()` signatures are unchanged: +Generation returns media and actual parameters together; see [generation metadata](generation_metadata.md): ```c sd_img_gen_params_t params; @@ -161,7 +165,10 @@ sd_img_gen_params_init(¶ms); /* Set prompt, original-resolution input images, and generation options. */ params.image_preprocess.rules = "target=init,mode=crop-resize,filter=lanczos;" "target=mask,filter=nearest-exact"; -bool ok = generate_image(ctx, ¶ms, &images, &count); +sd_image_generation_result_t* result = NULL; +bool ok = generate_image(ctx, ¶ms, &result); +/* Consume result->items and their generation_info on success. */ +free_sd_image_generation_result(result); ``` Both generation parameter initializers set `image_preprocess.rules` to `NULL`, diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index 04d9c4f73c..e8036c60ea 100644 --- a/docs/int8_convrot.md +++ b/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,14 @@ 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. +- 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 currently require a group size of `256`; other supported group sizes use CPU execution. ## Example diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 391850c2c2..060926a590 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -76,9 +76,9 @@ on a B200, logged once). cuDNN runs F16 with F32 accumulation: for the H3 DiT at takes about 8 ms per attention call on a B200 against 28 ms for the sage kernel and 43 ms for the ggml flash attention kernel, and it is closer to an exact (fp64) reference than either, so frames and audio differ slightly from the previous kernels. `GGML_CUDA_CUDNN_ATTN=0` restores them and -`GGML_CUDA_CUDNN_ATTN=1` uses cuDNN for the flash attention op on any Ampere or newer GPU; -`GGML_CUDA_CUDNN_SAGE=0` / `=1` does the same for the sage op only, and -`GGML_CUDA_CUDNN_ATTN_BF16=1` runs cuDNN in BF16 (less accurate, same speed). +`GGML_CUDA_CUDNN_ATTN=1` uses cuDNN for the flash attention op on any Ampere or newer GPU, and +`GGML_CUDA_CUDNN_ATTN_BF16=1` runs cuDNN in BF16 (less accurate, same speed). The sage op no +longer has a cuDNN path: SageAttention moved out of ggml upstream, so `GGML_CUDA_CUDNN_SAGE` is gone. The video VAE decodes one 16x16 latent tile (256x256 pixels) per decoder graph by default. `SD_H3_VAE_TILE=N` uses N x N latent tiles instead, overriding `--vae-tile-size` (20 decodes @@ -163,13 +163,76 @@ latents condition the diffusion transformer. Paired video and audio references share the same timeline. Ref2VA cannot be combined with `--init-img` or `--end-img` in one request. -Reference images keep their aspect ratio and are only downscaled when their -pixel area exceeds the requested generation canvas. +Reference images keep their aspect ratio and are only downscaled when they +exceed the configured limit; neither mode scales a reference up. The limit is +selected by `ref_image_size`. The default, `match`, uses the pixel area of the +requested generation canvas. Pass `--ref-image-args "ref_image_size=max"` (or +`ref_image_size` in `sd_vid_gen_params_t.ref_image_args`) to use a 2048 px short +edge instead, which preserves more identity detail at a higher sampling cost: +reference latents are attended at every sampling step. + +Whichever limit applies, each axis is then snapped to the nearest 32-pixel +multiple. An axis of at least 16 pixels can therefore grow by up to 16 pixels +past the scaled size (a 48-pixel axis becomes 64), and references are never +narrower than 32 pixels. +`resize_before_vae=false` is not supported here; MiniMax-H3 keeps references on +the 32-pixel grid and reports a warning. 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 ` + `--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. diff --git a/docs/model_config.md b/docs/model_config.md index dda93a837f..0d9290c1d5 100644 --- a/docs/model_config.md +++ b/docs/model_config.md @@ -1,7 +1,7 @@ # Model Configuration Conventions -This document describes the conventions for model configuration structs and -weight-based configuration detection. +This document describes model configuration structs, weight-based configuration +detection, and the definitions used to integrate models into the runtime. ## Config Types @@ -116,3 +116,113 @@ struct. If a model has no weight-derived config today, it may still provide `detect_from_weights` for API consistency, but it should not print a config detection log unless it actually derives values from weights. + +## Model Integration Definitions + +Every `SDVersion` has an explicit `ModelDefinition` in +`src/model/integrations/`. Related variants share an integration file and runner +factory. Network blocks and their weight-derived configuration remain under +`src/model/diffusion/`; reusable text encoder implementations remain under +`src/model/te/`. Weight-name conversion remains in `src/name_conversion.cpp`. + +To add a model: + +1. Add its identity to `SDVersion` without renumbering existing values. +2. Implement the network and `*Config::detect_from_weights` following the + conventions above, or reuse an existing runner. +3. Add a definition in its integration file, including its component + factory, traits, geometry, and sampling policy. +4. Add the family's explicit registration in `src/model_registry.cpp`. For a new + variant of a registered family, add its definition to that family's list. +5. Add version detection signatures and any required weight-name mappings in + `src/model_loader.cpp` and `src/name_conversion.cpp` respectively. Update shared + conditioning or runtime code only when existing behavior cannot serve the model. + +Use the existing family integration as the starting point for a variant. For a +new family, `src/model/integrations/moss_image.cpp` is a compact example of a +conditioner/diffusion factory and its definition. Definitions collect model +construction and defaults; they do not replace version detection, name conversion, +or model-specific inference algorithms. + +The registry checks for missing, duplicate, and invalid version registrations. +Explicit references keep integrations reachable in static-library builds. The +existing CMake glob includes integration sources automatically. No legacy +builder, display-name table, or prediction-default fallback needs updating. + +### Detection And Weight Conversion + +Version detection remains in `ModelLoader::get_sd_version()` in +`src/model_loader.cpp`. Add model signatures there and preserve the order of +overlapping checks. Detection uses tensor names and metadata, and can run before +and after name conversion, so signatures must work at both stages. Do not load +tensor data to identify an architecture. The registry looks up definitions by +the detected version; it does not perform detection. + +Weight-name mappings and their version-based dispatch remain centralized in +`src/name_conversion.cpp`. Add new mappings there when needed. Inference, GGUF +export, quantization, and LoRA conversion share the `convert_tensor_name` entry +point; preserve its prefix normalization and component conversion order. + +### Components And Conditioning + +The core factory receives `model_builders::Context`, the tokenizer configuration, +and the output runners. Backend checks, optional IP-Adapter setup, scale +overrides, and tokenizer usage validation remain in the common builder. A +factory returns failure if any required additional backend cannot initialize. + +LLM architecture selection, tokenizer requirements, prompt preparation, and +encoding remain centralized in `src/conditioning/conditioner.hpp`. Integration +factories instantiate the conditioner with the model version; models requiring +new LLM conditioning behavior still need changes in that shared implementation. + +`vae_kind` is the single source of truth for the VAE family; compatibility +helpers derive their results from it. `KL`, `FLUX`, and `FLUX2` share the +`AutoEncoderKL` implementation while retaining their distinct latent formats. +The related traits have separate responsibilities: + +- `supports_video_generation` describes the model's video generation capability. +- `video_tae` selects `TinyVideoAutoEncoder` instead of `TinyImageAutoEncoder`. + Image models can use a video TAE, so this is independent of video generation. +- `fake_vae` selects FakeVAE for pixel-space models. It is separate from + `vae_kind`, which also controls latent-format compatibility helpers. +- `vae_format_override` allows PiD's VAE format override. + +`configure_vae` handles model-specific KL adjustments such as +PixArt's scale and SDXL's convolution scale. VAE compatibility traits also drive +latent transformations and previews through the existing helper functions. + +### Sampling And Geometry + +`ModelSamplingPolicy` groups the defaults first: `default_prediction`, +`default_flow_shift`, `default_sample_method`, and `default_scheduler`. Set only +the values that differ from the struct defaults. `prediction_from_weights` and +`flow_shift_from_weights` preserve checkpoint-dependent defaults during automatic +prediction selection; `probe_v_prediction` enables SD2's runtime v-prediction +probe. Explicit prediction settings take precedence over automatic selection. +MOSS's default unconditional skip layer stays +in the diffusion engine, with active unconditional SLG taking precedence. + +The resolved `prediction_t` selects the denoiser directly. Flux3, MiniMax H3 AV, +and Ming-Image use `FLUX3_FLOW_PRED`, `H3_AV_FLOW_PRED`, and +`MING_IMAGE_FLOW_PRED`; their CLI overrides are `flux3_flow`, `h3_av_flow`, and +`ming_image_flow`. Explicit `sd3_flow` always selects `DiscreteFlowDenoiser`, +and `flux_flow` always selects `FluxFlowDenoiser`, regardless of model version. +Automatic selection retains each model's default implementation. New prediction +values are appended before `PREDICTION_COUNT`, preserving existing valid values; +clients using the automatic-selection sentinel must rebuild with the new header. + +`ModelGeometry` declares `diffusion_down_factor`, `latent_channels`, and +`image_channels`. `ref_image_preset` defaults to `"default"`; set it only for models +that need another reference-image preset. + +The common pipeline executes these policies; specialized video, multimodal, +cache, and extension algorithms remain in their existing runtime components. +Model-specific diffusion extra parameters are assembled in the diffusion engine. +Video frame conversion and alignment also remain in the diffusion engine. +`get_vae_scale_factor()` retains PiD's pixel-space special case and otherwise +uses the VAE's scale factor. + +Validate migrations against the previous implementation using synthetic metadata, +name conversions, prompt preparation, default/override selection, and focused +builds. Real generation checks are still needed to validate numerical behavior +when changing network or encoding algorithms. diff --git a/docs/p_h3_image.md b/docs/p_h3_image.md new file mode 100644 index 0000000000..20de6fb202 --- /dev/null +++ b/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 . Give the cat the scarf from , preserving its face and the background." -v --offload-to-cpu --rng cpu -H 2912 -W 2912 +``` + +Reference images are numbered in argument order as ``, ``, and so on. A plain edit instruction gets `Edit . ` prepended automatically. Prompts starting with `Edit ` 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. diff --git a/docs/pe_cache.md b/docs/pe_cache.md new file mode 100644 index 0000000000..e181959991 --- /dev/null +++ b/docs/pe_cache.md @@ -0,0 +1,60 @@ +# Position embedding cache + +`sd::PECache` stores CPU-side `sd::Tensor` data. Each model runner owns +its cache and calls the existing `make_input()` on every graph build. Cache +hits skip PE generation; allocation and upload remain ordinary graph-input +operations. `GGMLRunner` has no PE-specific code or resource hooks. + +The cache keeps the most recent key and tensor for each named PE input. A +matching key returns the existing tensor by reference. Changing its parameters +replaces that input's cached result. Different named inputs have independent +storage. Alternating two layouts under the same name recomputes their PE; +there is no LRU, device-buffer cache, memory budget, or graph lifecycle callback. +CPU data lives with the model runner, including across `runner_end()` calls, +and is released when its named entry is replaced or the cache is destroyed. + +## Adding a model + +Include `core/pe_cache.h` and declare a member: + +```cpp +sd::PECache pe_cache_; +``` + +Construct the key before generating PE and place the generation code in the +lazy callback: + +```cpp +sd::PECacheKey key("model.image_pe"); +key.add(height).add(width).add(head_dim).add(theta) + .add(circular_x_enabled).add(circular_y_enabled); +auto pe = make_input(pe_cache_.get_or_create(key, [&]() { + auto values = generate_position_embedding(/* existing arguments */); + return sd::Tensor({2, 2, head_dim / 2, token_count}, std::move(values)); +})); +``` + +Keys must cover every variable affecting PE values or shape: geometry, text +length, frequency settings, reference-image ordering and offsets, circular +settings, and video/audio positions or frame rates as applicable. +`add_shape()` records all four ggml dimensions, including trailing singleton +dimensions. `add_shapes()` records an ordered list of reference shapes. +`add(sd::Tensor)` records both shape and values for explicit position inputs. +Keys compare their complete encoded fields rather than a data pointer or hash. + +Use distinct names for PE inputs that coexist in a graph, such as Anima's +adapter query and key PE. `make_input()` defers reading the CPU data until graph +execution: do not replace a named entry while an unfinished graph references +it. Adding or updating other named entries does not invalidate its storage. +Return tensor shapes with the contiguous dimension first, and preserve each +model's existing F32 values, layout, and final transformations. + +## Current integrations + +The cache is used by MiniT2I, the Flux runner and its variants, Flux3, Qwen Image, +Qwen Image 2.1, Z-Image and L2P, Anima, Boogu, ERNIE Image, Fibo2, Ideogram4, +Iris, Krea2, Lens, LLaDA Image, MAGE Flow, Ming Image, MOSS Image, PiD, PixArt, +Wan, Hunyuan Video, LingBot Video, and LTX audio/video. + +Position embeddings stored as model weights and graph-side RoPE operations +remain on their existing paths, including MiniMax H3 and SenseNova U1. diff --git a/docs/performance.md b/docs/performance.md index 18491039eb..2e10c36e20 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -169,6 +169,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 diff --git a/docs/rpc.md b/docs/rpc.md index ea9edf1c31..cb28e256f0 100644 --- a/docs/rpc.md +++ b/docs/rpc.md @@ -131,7 +131,7 @@ If you want to allow connections from other machines on the network: > **Drivers & Hardware:** Ensure the Server machine has the necessary drivers installed and functional (e.g., Nvidia Drivers for CUDA, Vulkan SDK, or Metal). If no devices are found, the server will simply fallback to CPU usage. - +``` + +The options can be given in either order; no model is required. The client exits +with an error if RPC support is unavailable or a server cannot provide devices. -### Step B: Run with RPC device +### Step C: Run with RPC device If everything is working correctly, you can now run the client while offloading some or all of the work to the RPC server. @@ -217,4 +217,4 @@ The client will map these servers to sequential device IDs (e.g., RPC0 from the ## 6. Performance Considerations -RPC performance is heavily dependent on network bandwidth, as large weights and activations must be transferred back and forth over the network, especially for large models, or when using high resolutions. For best results, ensure your network connection is stable and has sufficient bandwidth (>1Gbps recommended). This shoumd not be a concern if you are running the server and client on the same machine, as the data transfer will happen over the loopback interface. \ No newline at end of file +RPC performance is heavily dependent on network bandwidth, as large weights and activations must be transferred back and forth over the network, especially for large models, or when using high resolutions. For best results, ensure your network connection is stable and has sufficient bandwidth (>1Gbps recommended). This shoumd not be a concern if you are running the server and client on the same machine, as the data transfer will happen over the loopback interface. diff --git a/docs/sage_attention.md b/docs/sage_attention.md index 46c854051e..2c5d9602d3 100644 --- a/docs/sage_attention.md +++ b/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 | | --- | --- | diff --git a/docs/taesd.md b/docs/taesd.md index 954086d180..8331f441bf 100644 --- a/docs/taesd.md +++ b/docs/taesd.md @@ -49,3 +49,5 @@ curl -L -o taeqi2_1.safetensors https://huggingface.co/madebyollin/taeqi2_1/reso ``` Replace `--vae PATH` with `--taesd taeqi2_1.safetensors` in the [Qwen Image 2.1 examples](qwen_image_2.1.md) to use it for encoding and decoding. For previews only, keep `--vae PATH` and add `--taesd taeqi2_1.safetensors --taesd-preview-only --preview tae`. + +`sd-server` accepts the same `--taesd` and `--taesd-preview-only` options; the preview method is selected per request instead of via `--preview`. diff --git a/docs/wan.md b/docs/wan.md index 87a960f305..2031de7df0 100644 --- a/docs/wan.md +++ b/docs/wan.md @@ -184,6 +184,8 @@ Notes: +TI2V 5B can decode multiple latent frames together when temporal tiling is enabled. For example, `--temporal-tiling --vae-tiling --vae-tile-size 256x256 --extra-tiling-args "temporal_tile_frames=4,temporal_decode_frames=2"` keeps up to four latent frames in each temporal tile and decodes up to two at once, with spatial tiles of 256×256 pixels. `temporal_decode_frames` defaults to `1`, preserving frame-by-frame decoding; other model versions ignore it. Increasing it requires more memory. Adjust the spatial tile size to fit the available memory. + ### Wan2.1 FLF2V 14B ``` diff --git a/examples/cli/CMakeLists.txt b/examples/cli/CMakeLists.txt index 6b1266cf58..90164cfc02 100644 --- a/examples/cli/CMakeLists.txt +++ b/examples/cli/CMakeLists.txt @@ -2,8 +2,10 @@ set(TARGET sd-cli) add_executable(${TARGET} ../common/common.cpp + ../common/generation_metadata.cpp ../common/log.cpp ../common/media_io.cpp + ../common/media_metadata.cpp image_metadata.cpp main.cpp ) diff --git a/examples/cli/README.md b/examples/cli/README.md index bf9de1862d..5789455491 100644 --- a/examples/cli/README.md +++ b/examples/cli/README.md @@ -31,14 +31,21 @@ rule for edge detection. See [Image preprocessing](../../docs/image_preprocessing.md) for input selectors, input defaults, downstream model processing, mask alignment, and examples. -Metadata mode inspects PNG/JPEG container metadata without loading any model: +Metadata mode inspects PNG, JPEG, WebP (including animated WebP), AVI, WebM, and Matroska +container metadata without loading any model. Use `--image` for both images and videos: ```bash ./bin/sd-cli -M metadata --image ./output.png ./bin/sd-cli -M metadata --image ./output.jpg --metadata-format json +./bin/sd-cli -M metadata --image ./output.webm +./bin/sd-cli -M metadata --image ./output.avi --metadata-format json ./bin/sd-cli -M metadata --image ./output.png --metadata-raw ./bin/sd-cli -M metadata --image ./output.png --metadata-all ``` For completely black or white images or videos, NaNs, and the `--linear-scale` / `--attn-scale` workaround, see [Troubleshooting](../../docs/troubleshooting.md). + +Generation outputs embed the core [generation record](../../docs/generation_metadata.md). +Images retain the native JSON and an additional WebUI-compatible infotext. +`--disable-image-metadata` omits both from files. diff --git a/examples/cli/image_metadata.cpp b/examples/cli/image_metadata.cpp index ba0d6da230..bb5ab7cf84 100644 --- a/examples/cli/image_metadata.cpp +++ b/examples/cli/image_metadata.cpp @@ -2,12 +2,14 @@ #include #include +#include #include #include #include #include #include #include +#include #include #include #include @@ -391,8 +393,39 @@ namespace { json& report, std::string& error); + json parse_exif_user_comment(const uint8_t* data, size_t size) { + if (size >= 8 && memcmp(data, "ASCII\0\0\0", 8) == 0) { + return trim_trailing_nuls(bytes_to_string(data + 8, data + size)); + } + if (size < 8 || memcmp(data, "UNICODE\0", 8) != 0) { + return json{{"raw_hex_preview", hex_preview(data, size)}}; + } + size_t offset = 8; + // WebUI's piexif writer uses big-endian UTF-16 without a BOM. + bool little_endian = false; + if (size >= 10 && ((data[8] == 0xfe && data[9] == 0xff) || (data[8] == 0xff && data[9] == 0xfe))) { + little_endian = data[8] == 0xff; + offset = 10; + } + if ((size - offset) % 2 != 0) { + return json{{"error", "truncated UTF-16 UserComment"}, {"raw_hex_preview", hex_preview(data, size)}}; + } + std::u16string text; + for (; offset < size; offset += 2) { + text.push_back(static_cast(little_endian ? data[offset] | (uint16_t(data[offset + 1]) << 8) + : (uint16_t(data[offset]) << 8) | data[offset + 1])); + } + try { + std::wstring_convert, char16_t> converter; + return trim_trailing_nuls(converter.to_bytes(text)); + } catch (const std::range_error&) { + return json{{"error", "invalid UTF-16 UserComment"}, {"raw_hex_preview", hex_preview(data, size)}}; + } + } + json parse_exif_value(const std::vector& data, size_t value_offset, + uint16_t tag, uint16_t type, uint32_t count, bool little_endian, @@ -466,8 +499,9 @@ namespace { break; } case 7: - value = bytes_to_string(data.data() + value_offset, - data.data() + value_offset + count); + value = tag == 0x9286 ? parse_exif_user_comment(data.data() + value_offset, count) + : json(bytes_to_string(data.data() + value_offset, + data.data() + value_offset + count)); break; case 9: { if (count == 1) { @@ -591,6 +625,7 @@ namespace { tag_json["value"] = parse_exif_value(data, value_offset, + tag, type, value_count, little_endian, @@ -1098,6 +1133,178 @@ namespace { return true; } + static bool parse_avi_chunks(const std::vector& data, size_t offset, size_t end, bool info, bool include_raw, json& result, std::string& error) { + while (offset < end) { + if (end - offset < 8) { + error = "truncated AVI chunk header"; + return false; + } + const auto name = bytes_to_string(data.data() + offset, data.data() + offset + 4); + const size_t size = read_u32_le(data, offset + 4); + offset += 8; + if (size > end - offset || (size & 1u) > end - offset - size) { + error = "AVI chunk exceeds its container"; + return false; + } + json entry = {{"entry_type", "chunk"}, {"name", name}, {"length", size}, {"metadata_like", info}}; + if (name == "LIST") { + if (size < 4) { + error = "truncated AVI LIST type"; + return false; + } + const auto type = bytes_to_string(data.data() + offset, data.data() + offset + 4); + entry["data"]["list_type"] = type; + if (type == "INFO" && !info && + !parse_avi_chunks(data, offset + 4, offset + size, true, include_raw, result, error)) + return false; + } else if (info) { + entry["data"]["text"] = trim_trailing_nuls(bytes_to_string(data.data() + offset, data.data() + offset + size)); + } else { + append_raw_preview(entry, data.data() + offset, size, include_raw); + } + result["entries"].push_back(std::move(entry)); + offset += size + (size & 1u); + } + return true; + } + + static bool parse_avi(const std::vector& data, bool include_raw, json& result, std::string& error) { + result = {{"format", "AVI"}, {"entries", json::array()}}; + size_t offset = 0; + while (offset < data.size()) { + if (data.size() - offset < 12 || memcmp(data.data() + offset, "RIFF", 4) != 0 || + memcmp(data.data() + offset + 8, offset == 0 ? "AVI " : "AVIX", 4) != 0) { + error = "invalid AVI RIFF header"; + return false; + } + const size_t size = read_u32_le(data, offset + 4); + if (size < 4 || size > data.size() - offset - 8 || + (size & 1u) > data.size() - offset - 8 - size) { + error = "AVI RIFF size exceeds file size"; + return false; + } + if (!parse_avi_chunks(data, offset + 12, offset + 8 + size, false, include_raw, result, error)) + return false; + offset += 8 + size + (size & 1u); + } + return true; + } + + static bool read_ebml_uint(const std::vector& data, size_t& offset, size_t end, bool id, uint64_t& value, bool& unknown) { + if (offset == end || data[offset] == 0) + return false; + uint8_t mask = 0x80; + size_t length = 1; + while (!(data[offset] & mask)) { + mask >>= 1; + ++length; + } + if (length > (id ? 4u : 8u) || length > end - offset) + return false; + value = data[offset++] & (id ? 0xff : mask - 1); + for (size_t i = 1; i < length; ++i) + value = (value << 8) | data[offset++]; + unknown = !id && value == (uint64_t(1) << (7 * length)) - 1; + return true; + } + + static bool is_webm_segment_child(uint64_t id) { + return id == 0x114D9B74 || id == 0x1549A966 || id == 0x1654AE6B || + id == 0x1F43B675 || id == 0x1C53BB6B || id == 0x1941A469 || + id == 0x1043A770 || id == 0x1254C367; + } + + static bool parse_webm_elements(const std::vector& data, size_t& offset, size_t end, uint64_t parent, bool unknown_cluster, int depth, bool include_raw, json& result, std::string& error) { + if (depth > 16) { + error = "WebM tag nesting exceeds limit"; + return false; + } + std::string tag_name; + std::string tag_text; + bool has_text = false; + while (offset < end) { + const auto start = offset; + uint64_t id, size; + bool unknown; + if (!read_ebml_uint(data, offset, end, true, id, unknown)) { + error = "invalid or truncated WebM element ID"; + return false; + } + // An unknown-size Cluster ends at the next Segment-level element. + if (unknown_cluster && is_webm_segment_child(id)) { + offset = start; + return true; + } + if (!read_ebml_uint(data, offset, end, false, size, unknown) || + (!unknown && size > end - offset)) { + error = "invalid or truncated WebM element size"; + return false; + } + if (unknown && id != 0x18538067 && id != 0x1F43B675) { + error = "unsupported unknown-size WebM element"; + return false; + } + const size_t stop = unknown ? end : offset + static_cast(size); + const bool container = (parent == 0 && (id == 0x1A45DFA3 || id == 0x18538067)) || + (parent == 0x18538067 && (id == 0x1254C367 || (id == 0x1F43B675 && unknown))) || + (parent == 0x1254C367 && id == 0x7373) || + ((parent == 0x7373 || parent == 0x67C8) && id == 0x67C8); + if (parent == 0 && id == 0x18538067) + result["has_segment"] = true; + if (parent == 0x1A45DFA3 && id == 0x4282) { + const auto type = bytes_to_string(data.data() + offset, data.data() + stop); + if (type != "webm" && type != "matroska") { + error = "EBML document is not WebM or Matroska"; + return false; + } + result["doc_type"] = type; + result["format"] = type == "webm" ? "WEBM" : "MATROSKA"; + } + if (parent == 0x67C8 && (id == 0x45A3 || id == 0x4487)) { + auto value = bytes_to_string(data.data() + offset, data.data() + stop); + if (id == 0x45A3) + tag_name = std::move(value); + else { + tag_text = std::move(value); + has_text = true; + } + } + if (parent == 0 || parent == 0x18538067) { + std::ostringstream name; + name << "0x" << std::hex << std::uppercase << id; + json entry = {{"entry_type", "element"}, {"name", name.str()}, {"length", stop - offset}, {"metadata_like", false}}; + if (unknown) + entry["unknown_size"] = true; + append_raw_preview(entry, data.data() + offset, stop - offset, include_raw); + result["entries"].push_back(std::move(entry)); + } + if (container) { + if (!parse_webm_elements(data, offset, stop, id, id == 0x1F43B675 && unknown, + depth + 1, include_raw, result, error)) + return false; + } else { + offset = stop; + } + } + if (parent == 0x67C8 && has_text) { + result["entries"].push_back({{"entry_type", "tag"}, {"name", tag_name}, {"metadata_like", true}, {"data", {{"text", tag_text}}}}); + } + return true; + } + + static bool parse_webm(const std::vector& data, bool include_raw, json& result, std::string& error) { + result = {{"format", "WEBM"}, {"entries", json::array()}}; + size_t offset = 0; + if (!parse_webm_elements(data, offset, data.size(), 0, false, 0, include_raw, result, error)) + return false; + if (!result.contains("doc_type") || !result.value("has_segment", false)) { + error = "missing WebM document type or Segment"; + return false; + } + result.erase("has_segment"); + return true; + } + std::string abbreviate(const std::string& value, bool brief) { if (!brief || value.size() <= 240) { return value; @@ -1116,7 +1323,7 @@ namespace { return; } if (value.is_primitive()) { - out << prefix << key << ": " << value.dump() << "\n"; + out << prefix << key << ": " << value.dump(-1, ' ', false, json::error_handler_t::replace) << "\n"; return; } if (value.is_array()) { @@ -1125,7 +1332,7 @@ namespace { if (item.is_string()) { out << prefix << " - " << abbreviate(item.get(), brief) << "\n"; } else if (item.is_primitive()) { - out << prefix << " - " << item.dump() << "\n"; + out << prefix << " - " << item.dump(-1, ' ', false, json::error_handler_t::replace) << "\n"; } else { out << prefix << " -\n"; for (auto it = item.begin(); it != item.end(); ++it) { @@ -1158,10 +1365,12 @@ namespace { } for (const auto& entry : entries) { - const bool is_chunk = entry.value("entry_type", "") == "chunk"; - const std::string name = is_chunk ? entry.value("name", "unknown") - : entry.value("marker", "unknown"); - out << (is_chunk ? "Chunk: " : "Segment: ") << name << "\n"; + const auto type = entry.value("entry_type", ""); + const auto name = entry.value("name", entry.value("marker", "unknown")); + const char* label = type == "chunk" ? "Chunk" : type == "tag" ? "Tag" + : type == "element" ? "Element" + : "Segment"; + out << label << ": " << name << "\n"; for (auto it = entry.begin(); it != entry.end(); ++it) { if (it.key() == "entry_type" || it.key() == "name" || it.key() == "marker" || it.key() == "metadata_like") { @@ -1211,7 +1420,13 @@ namespace { return parse_webp(data, include_raw, report, error); } - error = "unsupported image format; only PNG, JPEG, and WebP are supported"; + if (data.size() >= 12 && memcmp(data.data(), "RIFF", 4) == 0 && + memcmp(data.data() + 8, "AVI ", 4) == 0) + return parse_avi(data, include_raw, report, error); + if (data.size() >= 4 && read_u32_be(data, 0) == 0x1A45DFA3) + return parse_webm(data, include_raw, report, error); + + error = "unsupported media format; only PNG, JPEG, WebP, AVI, WebM, and Matroska are supported"; return false; } @@ -1230,7 +1445,7 @@ bool print_image_metadata(const std::string& image_path, if (options.output_format == MetadataOutputFormat::JSON) { visible_report["file"] = image_path; - out << visible_report.dump(2) << "\n"; + out << visible_report.dump(2, ' ', false, json::error_handler_t::replace) << "\n"; } else { print_text_report(image_path, visible_report, options.include_structural, options.brief, out); } diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index ca8ce4e79e..591a151123 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -47,7 +47,6 @@ struct SDCliParams { int preview_interval = 1; std::string preview_path = "preview.png"; int preview_fps = 16; - bool taesd_preview = false; bool preview_noisy = false; bool color = false; bool metadata_raw = false; @@ -70,7 +69,7 @@ struct SDCliParams { &output_path}, {"", "--image", - "path to the image to inspect (for metadata mode)", + "path to the image or video to inspect (for metadata mode)", 0, &image_path}, {"", @@ -114,10 +113,6 @@ struct SDCliParams { "--color", "colors the logging tags according to level", true, &color}, - {"", - "--taesd-preview-only", - std::string("prevents usage of taesd for decoding the final image. (for use with --preview ") + previews_str[PREVIEW_TAE] + ")", - true, &taesd_preview}, {"", "--preview-noisy", "enables previewing noisy inputs of the models rather than the denoised outputs", @@ -232,7 +227,7 @@ struct SDCliParams { } } else { if (image_path.empty()) { - LOG_ERROR("error: metadata mode needs an image path (--image)"); + LOG_ERROR("error: metadata mode needs an image or video path (--image)"); return false; } if (metadata_format != "text" && metadata_format != "json") { @@ -268,7 +263,6 @@ struct SDCliParams { << " preview_interval: " << preview_interval << ",\n" << " preview_path: \"" << preview_path << "\",\n" << " preview_fps: " << preview_fps << ",\n" - << " taesd_preview: " << (taesd_preview ? "true" : "false") << ",\n" << " preview_noisy: " << (preview_noisy ? "true" : "false") << ",\n" << " imatrix_out: \"" << imatrix_out << "\",\n" << " metadata_raw: " << (metadata_raw ? "true" : "false") << ",\n" @@ -301,6 +295,10 @@ void parse_args(int argc, const char** argv, SDCliParams& cli_params, SDContextP log_level = cli_params.log_level; log_color = cli_params.color; + if (ctx_params.list_devices) { + exit(ctx_params.print_devices() ? 0 : 1); + } + bool valid = cli_params.resolve_and_validate(); if (valid && cli_params.mode != METADATA) { valid = ctx_params.resolve_and_validate(cli_params.mode) && @@ -443,10 +441,11 @@ static fs::path get_video_audio_sidecar_path(const SDCliParams& cli_params) { bool save_results(const SDCliParams& cli_params, const SDContextParams& ctx_params, const SDGenerationParams& gen_params, - sd_image_t* results, + SDImageVec& results, int num_results, - const sd_audio_t* generated_audio = nullptr) { - if (results == nullptr || num_results <= 0) { + const sd_audio_t* generated_audio, + const std::string& video_info) { + if (results.empty() || num_results <= 0) { return false; } @@ -489,12 +488,10 @@ bool save_results(const SDCliParams& cli_params, if (!img.data) return false; - int images_per_batch = gen_params.batch_count > 0 ? std::max(1, num_results / gen_params.batch_count) : 1; - const int64_t metadata_seed = cli_params.mode == VID_GEN ? gen_params.seed : gen_params.seed + idx / images_per_batch; - std::string params = gen_params.embed_image_metadata - ? get_image_params(ctx_params, gen_params, metadata_seed, cli_params.mode) - : ""; - const bool ok = write_image_to_file(path.u8string(), img.data, img.width, img.height, img.channel, params, cli_params.compression_quality); + std::string params = gen_params.embed_image_metadata + ? results.generation_info(idx) + : ""; + const bool ok = write_image_to_file(path.u8string(), img.data, img.width, img.height, img.channel, params, cli_params.compression_quality); LOG_INFO("save result image %d to '%s' (%s)", idx, path.u8string().c_str(), ok ? "success" : "failure"); return ok; }; @@ -536,7 +533,7 @@ bool save_results(const SDCliParams& cli_params, if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm") ext = ".avi"; std::string params = gen_params.embed_image_metadata - ? get_image_params(ctx_params, gen_params, gen_params.seed, cli_params.mode) + ? video_info : ""; fs::path video_path = base_path; @@ -544,7 +541,7 @@ bool save_results(const SDCliParams& cli_params, std::string final_ext_lower = ext.u8string(); std::transform(final_ext_lower.begin(), final_ext_lower.end(), final_ext_lower.begin(), ::tolower); const bool mux_audio = generated_audio != nullptr && (final_ext_lower == ".avi" || final_ext_lower == ".webm"); - if (create_video_from_sd_images(video_path.u8string().c_str(), results, num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr, params) == 0) { + if (create_video_from_sd_images(video_path.u8string().c_str(), results.data(), num_results, gen_params.fps, cli_params.compression_quality, mux_audio ? generated_audio : nullptr, params) == 0) { LOG_INFO("save result video to '%s'", video_path.u8string().c_str()); if (generated_audio != nullptr && !mux_audio) { fs::path wav_path = video_path; @@ -609,30 +606,29 @@ static bool apply_adetailer(sd_ctx_t* sd_ctx, continue; } sd_img_gen_params_t ad_generation_params = img_gen_params; - ad_generation_params.seed = img_gen_params.seed + i; + ad_generation_params.seed = generation_seed(results.generation_info(i), img_gen_params.seed); if (mode == IMG_GEN) { ad_generation_params.width = 512; ad_generation_params.height = 512; ad_generation_params.strength = 0.4f; } - sd_image_t* detailed_images = nullptr; - int detailed_count = 0; + sd_image_generation_result_t* detailed = nullptr; if (!adetail_image(ad_ctx.get(), sd_ctx, results[i], &ad_params, &ad_generation_params, - &detailed_images, - &detailed_count) || - detailed_count <= 0 || detailed_images == nullptr || detailed_images[0].data == nullptr) { - free_sd_images(detailed_images, detailed_count); + &detailed) || + detailed == nullptr || detailed->count <= 0 || detailed->items[0].image.data == nullptr) { + free_sd_image_generation_result(detailed); LOG_ERROR("ADetailer failed for image %d", i + 1); return false; } + results.append_generation_info(i, detailed->items[0].generation_info, ctx_params.lora_model_dir); free(results[i].data); - results[i] = detailed_images[0]; - detailed_images[0] = {0, 0, 0, nullptr}; - free_sd_images(detailed_images, detailed_count); + results[i] = detailed->items[0].image; + detailed->items[0].image = {}; + free_sd_image_generation_result(detailed); } return true; } @@ -881,11 +877,12 @@ int main(int argc, const char* argv[]) { } } - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(cli_params.taesd_preview); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(); SDImageVec results; int num_results = 0; sd_audio_t* generated_audio = nullptr; + std::string video_info; if (cli_params.mode == UPSCALE) { num_results = 1; @@ -898,18 +895,6 @@ int main(int argc, const char* argv[]) { return 1; } - if (gen_params.sample_params.sample_method == SAMPLE_METHOD_COUNT) { - gen_params.sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get()); - } - - if (gen_params.high_noise_sample_params.sample_method == SAMPLE_METHOD_COUNT) { - gen_params.high_noise_sample_params.sample_method = sd_get_default_sample_method(sd_ctx.get()); - } - - if (gen_params.sample_params.scheduler == SCHEDULER_COUNT) { - gen_params.sample_params.scheduler = sd_get_default_scheduler(sd_ctx.get(), gen_params.sample_params.sample_method); - } - sd_img_gen_params_t img_gen_params{}; const bool use_img_gen_params = cli_params.mode == IMG_GEN || cli_params.mode == ADETAILER; if (use_img_gen_params) { @@ -917,23 +902,24 @@ int main(int argc, const char* argv[]) { } if (cli_params.mode == IMG_GEN) { - sd_image_t* generated_images = nullptr; - if (!generate_image(sd_ctx.get(), &img_gen_params, &generated_images, &num_results)) { - generated_images = nullptr; - num_results = 0; - } - results.adopt(generated_images, num_results); + sd_image_generation_result_t* generated = nullptr; + generate_image(sd_ctx.get(), &img_gen_params, &generated); + results.adopt(generated, ctx_params.lora_model_dir); + num_results = results.count(); } else if (cli_params.mode == ADETAILER) { num_results = 1; results.push_back(gen_params.init_image.release()); } else if (cli_params.mode == VID_GEN) { - sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); - sd_image_t* generated_video = nullptr; - if (!generate_video(sd_ctx.get(), &vid_gen_params, &generated_video, &num_results, &generated_audio, &cli_params.preview_fps)) { - generated_video = nullptr; + sd_vid_gen_params_t vid_gen_params = gen_params.to_sd_vid_gen_params_t(); + sd_video_generation_result_t* generated = nullptr; + if (generate_video(sd_ctx.get(), &vid_gen_params, &generated)) { + gen_params.fps = cli_params.preview_fps = generated->fps; + generated_audio = generated->audio; + generated->audio = nullptr; + num_results = generated->frame_count; + results.adopt_video(generated, ctx_params.lora_model_dir); + video_info = results.generation_info(0); } - gen_params.fps = cli_params.preview_fps; - results.adopt(generated_video, num_results); } if (!results) { @@ -973,21 +959,20 @@ int main(int argc, const char* argv[]) { SDImageOwner current_image(results[i]); results[i] = {0, 0, 0, nullptr}; for (int u = 0; u < gen_params.upscale_repeats; ++u) { - sd_image_t* upscaled_images = nullptr; - int upscaled_count = 0; - bool upscale_ok = upscale(upscaler_ctx.get(), - current_image.get(), - upscale_factor, - &upscaled_images, - &upscaled_count); - if (!upscale_ok || upscaled_count <= 0 || upscaled_images[0].data == nullptr) { - free_sd_images(upscaled_images, upscaled_count); + sd_image_generation_result_t* upscaled = nullptr; + bool upscale_ok = upscale(upscaler_ctx.get(), + current_image.get(), + upscale_factor, + &upscaled); + if (!upscale_ok || !upscaled || upscaled->count <= 0 || upscaled->items[0].image.data == nullptr) { + free_sd_image_generation_result(upscaled); LOG_ERROR("upscale failed"); break; } - sd_image_t upscaled_image = upscaled_images[0]; - upscaled_images[0] = {0, 0, 0, nullptr}; - free_sd_images(upscaled_images, upscaled_count); + results.append_generation_info(i, upscaled->items[0].generation_info); + sd_image_t upscaled_image = upscaled->items[0].image; + upscaled->items[0].image = {}; + free_sd_image_generation_result(upscaled); current_image.reset(upscaled_image); } results[i] = current_image.release(); // Set the final upscaled image as the result @@ -995,7 +980,9 @@ int main(int argc, const char* argv[]) { } } - if (!save_results(cli_params, ctx_params, gen_params, results.data(), num_results, generated_audio)) { + if (!video_info.empty() && results.count() > 0) + video_info = results.generation_info(0); + if (!save_results(cli_params, ctx_params, gen_params, results, num_results, generated_audio, video_info)) { free_sd_audio(generated_audio); return 1; } diff --git a/examples/common/common.cpp b/examples/common/common.cpp index b9fb780fdb..1b4996925d 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -464,7 +464,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}, {"", @@ -596,6 +596,11 @@ ArgOptions SDContextParams::get_options() { }; options.bool_options = { + {"", + "--list-devices", + "list available ggml backend devices, including --rpc-servers (one 'namedescription' per line), and exit; " + "the names are the device names accepted by --backend and --params-backend", + true, &list_devices}, {"", "--disable-prefetch", "disable asynchronous next-segment weight prefetch (defaults to false)", @@ -636,6 +641,10 @@ ArgOptions SDContextParams::get_options() { "--sage-attn", "use native CUDA SageAttention in the diffusion model, with flash/default attention fallback", true, &sage_attn}, + {"", + "--taesd-preview-only", + "prevents usage of taesd for decoding the final image; the VAE keeps encoding / decoding and taesd is only used for step previews", + true, &taesd_preview_only}, }; auto on_type_arg = [&](int argc, const char** argv, int index) { @@ -809,7 +818,7 @@ ArgOptions SDContextParams::get_options() { on_sampler_rng_arg}, {"", "--prediction", - "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow]", + "prediction type override, one of [eps, v, edm_v, sd3_flow, flux_flow, sefi_flow, minit2i_flow, sensenova_u1_flow, flux3_flow, h3_av_flow, ming_image_flow]", on_prediction_arg}, {"", "--lora-apply-mode", @@ -819,23 +828,22 @@ ArgOptions SDContextParams::get_options() { "but it usually offers faster inference speed and, in some cases, lower memory usage. " "The at_runtime mode, on the other hand, is exactly the opposite.", on_lora_apply_mode_arg}, - {"", - "--list-devices", - "list available ggml backend devices (one 'namedescription' per line) and exit; " - "the names are the device names accepted by --backend and --params-backend", - [](int /*argc*/, const char** /*argv*/, int /*index*/) { - size_t device_list_size = sd_list_devices(nullptr, 0); - std::vector devices(device_list_size + 1); - sd_list_devices(devices.data(), devices.size()); - fputs(devices.data(), stdout); - std::exit(0); - return 0; - }}, }; return options; } +bool SDContextParams::print_devices() const { + if (!sd_add_rpc_servers(rpc_servers.c_str())) { + return false; + } + size_t device_list_size = sd_list_devices(nullptr, 0); + std::vector devices(device_list_size + 1); + sd_list_devices(devices.data(), devices.size()); + fputs(devices.data(), stdout); + return true; +} + void SDContextParams::build_embedding_map() { static const std::vector valid_ext = {".gguf", ".safetensors", ".pt", ".ckpt"}; @@ -1010,6 +1018,7 @@ std::string SDContextParams::to_string() const { << " attn_scale: " << attn_scale << ",\n" << " diffusion_conv_direct: " << (diffusion_conv_direct ? "true" : "false") << ",\n" << " vae_conv_direct: " << (vae_conv_direct ? "true" : "false") << ",\n" + << " taesd_preview_only: " << (taesd_preview_only ? "true" : "false") << ",\n" << " prediction: " << sd_prediction_name(prediction) << ",\n" << " lora_apply_mode: " << sd_lora_apply_mode_name(lora_apply_mode) << ",\n" << " force_sdxl_vae_conv_scale: " << (force_sdxl_vae_conv_scale ? "true" : "false") << "\n" @@ -1017,7 +1026,7 @@ std::string SDContextParams::to_string() const { return oss.str(); } -sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { +sd_ctx_params_t SDContextParams::to_sd_ctx_params_t() { prepare_backend_assignments(); embedding_vec.clear(); embedding_vec.reserve(embedding_map.size()); @@ -1067,7 +1076,7 @@ sd_ctx_params_t SDContextParams::to_sd_ctx_params_t(bool taesd_preview) { sd_ctx_params.sage_attn = sage_attn; sd_ctx_params.linear_scale = linear_scale; sd_ctx_params.attn_scale = attn_scale; - sd_ctx_params.tae_preview_only = taesd_preview; + sd_ctx_params.tae_preview_only = taesd_preview_only; sd_ctx_params.diffusion_conv_direct = diffusion_conv_direct; sd_ctx_params.vae_conv_direct = vae_conv_direct; sd_ctx_params.force_sdxl_vae_conv_scale = force_sdxl_vae_conv_scale; @@ -1185,7 +1194,7 @@ ArgOptions SDGenerationParams::get_options() { &extra_sample_args}, {"", "--extra-tiling-args", - "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1)", + "extra VAE tiling args, key=value list. Supported video VAEs accept temporal_tile_frames/temporal_tile_size (default: 4), temporal_tile_overlap (default: 1); Wan2.2 TI2V 5B accepts temporal_decode_frames (default: 1)", (int)',', &extra_tiling_args}, {"", @@ -1521,6 +1530,8 @@ 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"; } if (!sigmas_str.empty() && sigmas_str.front() == '[') { sigmas_str.erase(0, 1); @@ -1768,11 +1779,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.", 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.", on_hires_sigmas_arg}, {"", "--skip-layers", @@ -1950,6 +1961,7 @@ static bool parse_image_json_field(const json& parent, return decode_base64_image(parent.at(key).get(), channels, 0, 0, out_image); } +/** Parses a JSON array of base64 image strings into owned images. */ static bool parse_image_array_json_field(const json& parent, const char* key, int channels, @@ -1979,6 +1991,84 @@ static bool parse_image_array_json_field(const json& parent, return true; } +/** Parses a JSON array of base64 WAV strings into owned audio clips. */ +static bool parse_audio_array_json_field(const json& parent, + const char* key, + std::vector& out_audios) { + if (!parent.contains(key)) { + return true; + } + if (parent.at(key).is_null()) { + out_audios.clear(); + return true; + } + if (!parent.at(key).is_array()) { + return false; + } + + out_audios.clear(); + for (const auto& item : parent.at(key)) { + if (!item.is_string()) { + return false; + } + std::string encoded = item.get(); + auto comma_pos = encoded.find(','); + if (comma_pos != std::string::npos) { + encoded = encoded.substr(comma_pos + 1); + } + std::vector audio_bytes = decode_base64_bytes(encoded); + if (audio_bytes.empty()) { + return false; + } + std::vector samples; + uint32_t sample_rate = 0; + uint32_t channels = 0; + if (!load_wav_from_memory(audio_bytes.data(), audio_bytes.size(), samples, sample_rate, channels)) { + return false; + } + SDAudioOwner audio; + audio.reset(std::move(samples), sample_rate, channels); + out_audios.push_back(std::move(audio)); + } + return true; +} + +/** Parses a JSON array of per-video base64 frame arrays into owned videos. */ +static bool parse_video_array_json_field(const json& parent, + const char* key, + std::vector>& out_videos) { + if (!parent.contains(key)) { + return true; + } + if (parent.at(key).is_null()) { + out_videos.clear(); + return true; + } + if (!parent.at(key).is_array()) { + return false; + } + + out_videos.clear(); + for (const auto& video : parent.at(key)) { + if (!video.is_array() || video.empty()) { + return false; + } + std::vector frames; + for (const auto& frame : video) { + if (!frame.is_string()) { + return false; + } + SDImageOwner image; + if (!decode_base64_image(frame.get(), 0, 0, 0, image)) { + return false; + } + frames.push_back(std::move(image)); + } + out_videos.push_back(std::move(frames)); + } + return true; +} + static bool parse_lora_json_field(const json& parent, const std::function& lora_path_resolver, std::map& lora_map, @@ -2093,6 +2183,7 @@ bool SDGenerationParams::parse_image_preprocess_json(const std::string& json_str return true; } +/** Populates generation params from a JSON payload (server APIs); returns false on any malformed field. */ bool SDGenerationParams::from_json_str( const std::string& json_str, const std::function& lora_path_resolver) { @@ -2219,7 +2310,7 @@ bool SDGenerationParams::from_json_str( if (target_custom_sigmas != nullptr && sample_json.contains("custom_sigmas") && sample_json["custom_sigmas"].is_array() && - !sample_json["custom_sigmas"].empty()) { // empty [] = "not specified", keep server/CLI default + !sample_json["custom_sigmas"].empty()) { // empty [] = "not specified", keep server/CLI default *target_custom_sigmas = sample_json["custom_sigmas"].get>(); } if (sample_json.contains("sample_method") && sample_json["sample_method"].is_string()) { @@ -2321,8 +2412,11 @@ bool SDGenerationParams::from_json_str( return false; } if (!parse_image_array_json_field(j, "ref_images", 0, ref_images) || - !parse_image_array_json_field(j, "control_frames", 3, control_frames)) { - LOG_ERROR("invalid input image array"); + !parse_image_array_json_field(j, "control_frames", 3, control_frames) || + !parse_audio_array_json_field(j, "ref_audios", ref_audios) || + !parse_audio_array_json_field(j, "ref_video_audios", ref_video_audios) || + !parse_video_array_json_field(j, "ref_videos", ref_videos)) { + LOG_ERROR("invalid input image/audio/video array"); return false; } @@ -2539,11 +2633,6 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s return false; } - if (seed < 0) { - srand((int)time(nullptr)); - seed = rand(); - } - if (strict) { batch_count = std::clamp(batch_count, 1, 8); sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100); @@ -2579,6 +2668,7 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s return true; } +/** Validates params for a generation mode; enforces mode-specific input rules. */ bool SDGenerationParams::validate(SDMode mode) { if (!image_preprocess.empty() && mode != IMG_GEN && mode != VID_GEN) { LOG_ERROR("--image-preprocess requires img_gen or vid_gen mode"); @@ -2632,12 +2722,13 @@ bool SDGenerationParams::validate(SDMode mode) { return false; } - if (ref_video_audio_paths.size() > ref_video_paths.size()) { + if (ref_video_audio_paths.size() + ref_video_audios.size() > ref_video_paths.size() + ref_videos.size()) { LOG_ERROR("error: each --ref-video-audio needs a corresponding --ref-video"); return false; } - if (mode != VID_GEN && (!ref_video_paths.empty() || !ref_video_audio_paths.empty() || !ref_audio_paths.empty())) { + if (mode != VID_GEN && (!ref_video_paths.empty() || !ref_video_audio_paths.empty() || !ref_audio_paths.empty() || + !ref_videos.empty() || !ref_video_audios.empty() || !ref_audios.empty())) { LOG_ERROR("error: reference video and audio inputs require vid_gen mode"); return false; } @@ -2910,6 +3001,7 @@ sd_vid_gen_params_t SDGenerationParams::to_sd_vid_gen_params_t() { params.circular_x = circular || circular_x; params.circular_y = circular || circular_y; params.image_preprocess = {image_preprocess.c_str()}; + params.ref_image_args = ref_image_args.c_str(); return params; } @@ -3019,301 +3111,3 @@ std::string SDGenerationParams::to_string() const { std::string version_string() { return std::string("stable-diffusion.cpp version ") + sd_version() + ", commit " + sd_commit(); } - -static std::string safe_json_string(const char* value) { - return value ? value : ""; -} - -static void set_json_basename_if_not_empty(json& target, const char* key, const std::string& path) { - if (!path.empty()) { - target[key] = sd_basename(path); - } -} - -static json build_sampling_metadata_json(const sd_sample_params_t& sample_params, - const std::vector& skip_layers, - const std::vector* custom_sigmas = nullptr) { - json sampling = { - {"steps", sample_params.sample_steps}, - {"eta", sample_params.eta}, - {"shifted_timestep", sample_params.shifted_timestep}, - {"flow_shift", sample_params.flow_shift}, - {"extra_sample_args", safe_json_string(sample_params.extra_sample_args)}, - {"guidance", - { - {"txt_cfg", sample_params.guidance.txt_cfg}, - {"img_cfg", sample_params.guidance.img_cfg}, - {"distilled_guidance", sample_params.guidance.distilled_guidance}, - {"slg", - { - {"scale", sample_params.guidance.slg.scale}, - {"layers", skip_layers}, - {"start", sample_params.guidance.slg.layer_start}, - {"end", sample_params.guidance.slg.layer_end}, - }}, - }}, - }; - if (sample_params.sample_method != SAMPLE_METHOD_COUNT) { - sampling["method"] = safe_json_string(sd_sample_method_name(sample_params.sample_method)); - } - if (sample_params.scheduler != SCHEDULER_COUNT) { - sampling["scheduler"] = safe_json_string(sd_scheduler_name(sample_params.scheduler)); - } - if (custom_sigmas != nullptr) { - sampling["custom_sigmas"] = *custom_sigmas; - } - return sampling; -} - -std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode) { - json root; - root["schema"] = "sdcpp.image.params/v1"; - root["mode"] = "img_gen"; - if (mode == VID_GEN) { - root["mode"] = "vid_gen"; - } else if (mode == ADETAILER) { - root["mode"] = "adetailer"; - } - root["generator"] = { - {"name", "stable-diffusion.cpp"}, - {"version", safe_json_string(sd_version())}, - {"commit", safe_json_string(sd_commit())}, - }; - root["seed"] = seed; - root["width"] = gen_params.get_resolved_width(); - root["height"] = gen_params.get_resolved_height(); - - root["prompt"] = { - {"positive", gen_params.prompt}, - {"negative", gen_params.negative_prompt}, - }; - if (!gen_params.ad_model_path.empty()) { - root["adetailer"] = { - {"model", sd_basename(gen_params.ad_model_path)}, - {"prompt", gen_params.ad_prompt}, - {"negative_prompt", gen_params.ad_negative_prompt}, - {"extra_args", gen_params.extra_ad_args}, - }; - } - root["sampling"] = build_sampling_metadata_json(gen_params.sample_params, - gen_params.skip_layers, - &gen_params.custom_sigmas); - - json models; - set_json_basename_if_not_empty(models, "model", ctx_params.model_path); - set_json_basename_if_not_empty(models, "clip_l", ctx_params.clip_l_path); - set_json_basename_if_not_empty(models, "clip_g", ctx_params.clip_g_path); - set_json_basename_if_not_empty(models, "clip_vision", ctx_params.clip_vision_path); - set_json_basename_if_not_empty(models, "t5xxl", ctx_params.t5xxl_path); - set_json_basename_if_not_empty(models, "llm", ctx_params.llm_path); - set_json_basename_if_not_empty(models, "llm_vision", ctx_params.llm_vision_path); - set_json_basename_if_not_empty(models, "diffusion_model", ctx_params.diffusion_model_path); - set_json_basename_if_not_empty(models, "high_noise_diffusion_model", ctx_params.high_noise_diffusion_model_path); - set_json_basename_if_not_empty(models, "uncond_diffusion_model", ctx_params.uncond_diffusion_model_path); - set_json_basename_if_not_empty(models, "vae", ctx_params.vae_path); - set_json_basename_if_not_empty(models, "taesd", ctx_params.taesd_path); - set_json_basename_if_not_empty(models, "control_net", ctx_params.control_net_path); - root["models"] = std::move(models); - - root["clip_skip"] = gen_params.clip_skip; - root["strength"] = gen_params.strength; - root["control_strength"] = gen_params.control_strength; - root["ip_adapter_strength"] = gen_params.ip_adapter_strength; - root["ref_image_args"] = gen_params.ref_image_args; - root["image_preprocess"] = gen_params.image_preprocess; - root["increase_ref_index"] = gen_params.increase_ref_index; - if (mode == VID_GEN) { - root["video"] = { - {"frame_count", gen_params.video_frames}, - {"fps", gen_params.fps}, - }; - root["moe_boundary"] = gen_params.moe_boundary; - root["vace_strength"] = gen_params.vace_strength; - root["high_noise_sampling"] = build_sampling_metadata_json(gen_params.high_noise_sample_params, - gen_params.high_noise_skip_layers); - } - - root["rng"] = safe_json_string(sd_rng_type_name(ctx_params.rng_type)); - if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) { - root["sampler_rng"] = safe_json_string(sd_rng_type_name(ctx_params.sampler_rng_type)); - } - - json loras = json::array(); - for (const auto& entry : gen_params.lora_map) { - loras.push_back({ - {"name", sd_basename(entry.first)}, - {"multiplier", entry.second}, - {"is_high_noise", false}, - }); - } - for (const auto& entry : gen_params.high_noise_lora_map) { - loras.push_back({ - {"name", sd_basename(entry.first)}, - {"multiplier", entry.second}, - {"is_high_noise", true}, - }); - } - if (!loras.empty()) { - root["loras"] = std::move(loras); - } - - if (gen_params.hires_enabled) { - root["hires"] = { - {"enabled", gen_params.hires_enabled}, - {"upscaler", gen_params.hires_upscaler}, - {"model", gen_params.hires_upscaler_model_path.empty() ? "" : sd_basename(gen_params.hires_upscaler_model_path)}, - {"scale", gen_params.hires_scale}, - {"target_width", gen_params.hires_width}, - {"target_height", gen_params.hires_height}, - {"steps", gen_params.hires_steps}, - {"denoising_strength", gen_params.hires_denoising_strength}, - {"custom_sigmas", gen_params.hires_custom_sigmas}, - {"upscale_tile_size", gen_params.hires_upscale_tile_size}, - }; - } - - if (gen_params.cache_params.mode != SD_CACHE_DISABLED) { - root["cache"] = { - {"requested_mode", gen_params.cache_mode}, - {"requested_option", gen_params.cache_option}, - {"mode", gen_params.cache_params.mode}, - {"scm_mask", gen_params.scm_mask}, - {"scm_policy_dynamic", gen_params.scm_policy_dynamic}, - {"reuse_threshold", gen_params.cache_params.reuse_threshold}, - {"start_percent", gen_params.cache_params.start_percent}, - {"end_percent", gen_params.cache_params.end_percent}, - {"error_decay_rate", gen_params.cache_params.error_decay_rate}, - {"use_relative_threshold", gen_params.cache_params.use_relative_threshold}, - {"reset_error_on_compute", gen_params.cache_params.reset_error_on_compute}, - {"Fn_compute_blocks", gen_params.cache_params.Fn_compute_blocks}, - {"Bn_compute_blocks", gen_params.cache_params.Bn_compute_blocks}, - {"residual_diff_threshold", gen_params.cache_params.residual_diff_threshold}, - {"max_warmup_steps", gen_params.cache_params.max_warmup_steps}, - {"max_cached_steps", gen_params.cache_params.max_cached_steps}, - {"max_continuous_cached_steps", gen_params.cache_params.max_continuous_cached_steps}, - {"taylorseer_n_derivatives", gen_params.cache_params.taylorseer_n_derivatives}, - {"taylorseer_skip_interval", gen_params.cache_params.taylorseer_skip_interval}, - {"spectrum_w", gen_params.cache_params.spectrum_w}, - {"spectrum_m", gen_params.cache_params.spectrum_m}, - {"spectrum_lam", gen_params.cache_params.spectrum_lam}, - {"spectrum_window_size", gen_params.cache_params.spectrum_window_size}, - {"spectrum_flex_window", gen_params.cache_params.spectrum_flex_window}, - {"spectrum_warmup_steps", gen_params.cache_params.spectrum_warmup_steps}, - {"spectrum_stop_percent", gen_params.cache_params.spectrum_stop_percent}, - }; - } - - if (gen_params.vae_tiling_params.enabled || - gen_params.vae_tiling_params.temporal_tiling || - !gen_params.extra_tiling_args.empty()) { - root["vae_tiling"] = { - {"enabled", gen_params.vae_tiling_params.enabled}, - {"temporal_tiling", gen_params.vae_tiling_params.temporal_tiling}, - {"tile_size_w", gen_params.vae_tiling_params.tile_size_w}, - {"tile_size_h", gen_params.vae_tiling_params.tile_size_h}, - {"target_overlap", gen_params.vae_tiling_params.target_overlap}, - {"rel_size_w", gen_params.vae_tiling_params.rel_size_w}, - {"rel_size_h", gen_params.vae_tiling_params.rel_size_h}, - {"extra_tiling_args", gen_params.extra_tiling_args}, - }; - } - - return root.dump(); -} - -std::string get_image_params(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode) { - std::string parameter_string; - if (gen_params.prompt_with_lora.size() != 0) { - parameter_string += gen_params.prompt_with_lora + "\n"; - } else { - parameter_string += gen_params.prompt + "\n"; - } - if (gen_params.negative_prompt.size() != 0) { - parameter_string += "Negative prompt: " + gen_params.negative_prompt + "\n"; - } - parameter_string += "Steps: " + std::to_string(gen_params.sample_params.sample_steps) + ", "; - parameter_string += "CFG scale: " + std::to_string(gen_params.sample_params.guidance.txt_cfg) + ", "; - if (gen_params.sample_params.guidance.slg.scale != 0 && gen_params.skip_layers.size() != 0) { - parameter_string += "SLG scale: " + std::to_string(gen_params.sample_params.guidance.slg.scale) + ", "; - parameter_string += "Skip layers: ["; - for (const auto& layer : gen_params.skip_layers) { - parameter_string += std::to_string(layer) + ", "; - } - parameter_string += "], "; - parameter_string += "Skip layer start: " + std::to_string(gen_params.sample_params.guidance.slg.layer_start) + ", "; - parameter_string += "Skip layer end: " + std::to_string(gen_params.sample_params.guidance.slg.layer_end) + ", "; - } - parameter_string += "Guidance: " + std::to_string(gen_params.sample_params.guidance.distilled_guidance) + ", "; - parameter_string += "Eta: " + std::to_string(gen_params.sample_params.eta) + ", "; - if (!gen_params.extra_sample_args.empty()) { - parameter_string += "Extra sample args: " + gen_params.extra_sample_args + ", "; - } - if (!gen_params.ad_model_path.empty()) { - parameter_string += "ADetailer model: " + sd_basename(gen_params.ad_model_path) + ", "; - if (!gen_params.ad_prompt.empty()) { - parameter_string += "ADetailer prompt: " + gen_params.ad_prompt + ", "; - } - if (!gen_params.ad_negative_prompt.empty()) { - parameter_string += "ADetailer negative prompt: " + gen_params.ad_negative_prompt + ", "; - } - if (!gen_params.extra_ad_args.empty()) { - parameter_string += "ADetailer args: " + gen_params.extra_ad_args + ", "; - } - } - parameter_string += "Seed: " + std::to_string(seed) + ", "; - parameter_string += "Size: " + std::to_string(gen_params.get_resolved_width()) + "x" + std::to_string(gen_params.get_resolved_height()) + ", "; - parameter_string += "Model: " + sd_basename(ctx_params.model_path) + ", "; - parameter_string += "RNG: " + std::string(sd_rng_type_name(ctx_params.rng_type)) + ", "; - if (ctx_params.sampler_rng_type != RNG_TYPE_COUNT) { - parameter_string += "Sampler RNG: " + std::string(sd_rng_type_name(ctx_params.sampler_rng_type)) + ", "; - } - parameter_string += "Sampler: " + std::string(sd_sample_method_name(gen_params.sample_params.sample_method)); - if (!gen_params.custom_sigmas.empty()) { - parameter_string += ", Custom Sigmas: ["; - for (size_t i = 0; i < gen_params.custom_sigmas.size(); ++i) { - std::ostringstream oss; - oss << std::fixed << std::setprecision(4) << gen_params.custom_sigmas[i]; - parameter_string += oss.str() + (i == gen_params.custom_sigmas.size() - 1 ? "" : ", "); - } - parameter_string += "]"; - } else if (gen_params.sample_params.scheduler != SCHEDULER_COUNT) { // Only show schedule if not using custom sigmas - parameter_string += " " + std::string(sd_scheduler_name(gen_params.sample_params.scheduler)); - } - parameter_string += ", "; - for (const auto& te : {ctx_params.clip_l_path, ctx_params.clip_g_path, ctx_params.t5xxl_path, ctx_params.llm_path, ctx_params.llm_vision_path}) { - if (!te.empty()) { - parameter_string += "TE: " + sd_basename(te) + ", "; - } - } - if (!ctx_params.diffusion_model_path.empty()) { - parameter_string += "Unet: " + sd_basename(ctx_params.diffusion_model_path) + ", "; - } - if (!ctx_params.uncond_diffusion_model_path.empty()) { - parameter_string += "Uncond Unet: " + sd_basename(ctx_params.uncond_diffusion_model_path) + ", "; - } - if (!ctx_params.vae_path.empty()) { - parameter_string += "VAE: " + sd_basename(ctx_params.vae_path) + ", "; - } - if (gen_params.clip_skip != -1) { - parameter_string += "Clip skip: " + std::to_string(gen_params.clip_skip) + ", "; - } - if (gen_params.hires_enabled) { - parameter_string += "Hires upscale: " + gen_params.hires_upscaler + ", "; - parameter_string += "Hires scale: " + std::to_string(gen_params.hires_scale) + ", "; - parameter_string += "Hires resize: " + std::to_string(gen_params.hires_width) + "x" + std::to_string(gen_params.hires_height) + ", "; - parameter_string += "Hires steps: " + std::to_string(gen_params.hires_steps) + ", "; - parameter_string += "Denoising strength: " + std::to_string(gen_params.hires_denoising_strength) + ", "; - if (!gen_params.hires_custom_sigmas.empty()) { - parameter_string += "Hires custom sigmas: " + vec_to_string(gen_params.hires_custom_sigmas) + ", "; - } - } - parameter_string += "Version: stable-diffusion.cpp"; - parameter_string += ", SDCPP: " + build_sdcpp_image_metadata_json(ctx_params, gen_params, seed, mode); - return parameter_string; -} diff --git a/examples/common/common.h b/examples/common/common.h index 91e75c6851..39d6011afc 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -1,6 +1,8 @@ #ifndef __EXAMPLES_COMMON_COMMON_H__ #define __EXAMPLES_COMMON_COMMON_H__ +#include "generation_metadata.h" + #include #include #include @@ -168,6 +170,7 @@ struct SDContextParams { std::string model_args; bool auto_fit = true; std::string rpc_servers; + bool list_devices = false; std::string effective_backend; std::string effective_params_backend; bool enable_mmap = false; @@ -179,6 +182,7 @@ struct SDContextParams { bool sage_attn = false; bool diffusion_conv_direct = false; bool vae_conv_direct = false; + bool taesd_preview_only = false; prediction_t prediction = PREDICTION_COUNT; lora_apply_mode_t lora_apply_mode = LORA_APPLY_AUTO; @@ -189,13 +193,14 @@ struct SDContextParams { float flow_shift = INFINITY; ArgOptions get_options(); + bool print_devices() const; void build_embedding_map(); void prepare_backend_assignments(); bool resolve(SDMode mode); bool validate(SDMode mode); bool resolve_and_validate(SDMode mode); std::string to_string() const; - sd_ctx_params_t to_sd_ctx_params_t(bool taesd_preview); + sd_ctx_params_t to_sd_ctx_params_t(); }; struct SDGenerationParams { @@ -335,13 +340,4 @@ struct SDGenerationParams { }; std::string version_string(); -std::string build_sdcpp_image_metadata_json(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode = IMG_GEN); -std::string get_image_params(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - int64_t seed, - SDMode mode = IMG_GEN); - #endif // __EXAMPLES_COMMON_COMMON_H__ diff --git a/examples/common/generation_metadata.cpp b/examples/common/generation_metadata.cpp new file mode 100644 index 0000000000..989e914462 --- /dev/null +++ b/examples/common/generation_metadata.cpp @@ -0,0 +1,155 @@ +#include "generation_metadata.h" + +#include +#include + +#include "json.hpp" + +static void relativize_lora_paths(nlohmann::json& info, const std::filesystem::path& root) { + if (info.contains("loras")) { + for (auto& lora : info["loras"]) { + const auto path = lora.value("path", ""); + if (path.empty()) + continue; + std::error_code error; + const auto absolute = std::filesystem::absolute(std::filesystem::u8path(path), error); + if (error) + continue; + const auto relative = absolute.lexically_normal().lexically_relative(root); + if (!relative.empty()) + lora["path"] = relative.generic_u8string(); + } + } + if (info.contains("inpaint_params")) + relativize_lora_paths(info["inpaint_params"], root); + if (info.contains("postprocess")) { + for (auto& operation : info["postprocess"]) + relativize_lora_paths(operation, root); + } +} + +std::string generation_info_with_relative_lora_paths(const std::string& info, const std::string& lora_model_dir) { + if (info.empty() || lora_model_dir.empty()) + return info; + std::error_code error; + const auto root = std::filesystem::absolute(std::filesystem::u8path(lora_model_dir), error); + if (error) + return info; + auto result = nlohmann::json::parse(info); + relativize_lora_paths(result, root.lexically_normal()); + return result.dump(); +} + +std::string merge_generation_info(const std::string& source, const std::string& operation) { + if (source.empty()) + return operation; + auto result = nlohmann::json::parse(source); + const auto next = nlohmann::json::parse(operation); + result["postprocess"].push_back(next); + return result.dump(); +} + +int64_t generation_seed(const std::string& info, int64_t fallback) { + if (info.empty()) + return fallback; + const auto value = nlohmann::json::parse(info); + const auto& params = value.contains("inpaint_params") ? value.at("inpaint_params") : value; + return params.contains("seed") ? std::stoll(params.at("seed").get()) : fallback; +} + +std::string generation_info_to_webui(const std::string& text) { + using nlohmann::json; + if (text.empty()) + return ""; + const auto info = json::parse(text); + if (info.value("schema", "") != "sdcpp.generation/v1") + return ""; + // Standalone postprocessing has no source diffusion settings to import. + if (!info.contains("sample_params")) + return ""; + const auto& sampling = info.at("sample_params"); + std::string result = info.value("prompt", ""); + if (info.contains("loras")) { + 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)) + continue; + auto name = lora.at("path").get(); + const auto separator = name.find_last_of("/\\"); + if (separator != std::string::npos) + name.erase(0, separator + 1); + const auto extension = name.find_last_of('.'); + if (extension != std::string::npos && extension > 0) + name.erase(extension); + if (name.empty() || name.find_first_of(":<>\r\n") != std::string::npos) + continue; + if (!result.empty()) + result += " "; + result += ""; + } + } + const auto negative = info.value("negative_prompt", ""); + if (!negative.empty()) + result += "\nNegative prompt: " + negative; + result += "\n"; + bool first = true; + const auto field = [&](const char* key, const json& value) { + std::string v = value.is_string() ? value.get() : value.dump(); + if (v.find_first_of(",\n\r:\"") != std::string::npos) + v = json(v).dump(-1, ' ', false); + if (!first) + result += ", "; + first = false; + result += std::string(key) + ": " + v; + }; + field("Steps", sampling.contains("custom_sigmas") ? json(sampling.at("custom_sigmas").size() - 1) : sampling.at("sample_steps")); + static const std::map samplers = { + {"euler", "Euler"}, {"euler_a", "Euler a"}, {"heun", "Heun"}, {"dpm2", "DPM2"}, {"dpm++2s_a", "DPM++ 2S a"}, {"dpm++2m", "DPM++ 2M"}, {"dpm++2m_sde", "DPM++ 2M SDE"}, {"lcm", "LCM"}}; + const auto method = sampling.at("sample_method").get(); + auto mapped = samplers.find(method); + field(mapped != samplers.end() ? "Sampler" : "SDCPP Sampler", mapped != samplers.end() ? mapped->second : method); + static const std::map schedules = { + {"discrete", "Normal"}, {"karras", "Karras"}, {"exponential", "Exponential"}, {"sgm_uniform", "SGM Uniform"}, {"simple", "Simple"}, {"ddim_uniform", "DDIM"}, {"beta", "Beta"}}; + if (sampling.contains("scheduler")) { + const auto scheduler = sampling.at("scheduler").get(); + auto mapped_schedule = schedules.find(scheduler); + field(mapped_schedule != schedules.end() ? "Schedule type" : "SDCPP Schedule", mapped_schedule != schedules.end() ? mapped_schedule->second : scheduler); + } + field("CFG scale", sampling.at("guidance").at("txt_cfg")); + field("Seed", info.at("seed")); + field("Size", std::to_string(info.at("width").get()) + "x" + std::to_string(info.at("height").get())); + if (info.contains("models")) { + const auto& models = info.at("models"); + auto model = models.value("diffusion_model", models.value("model", "")); + const auto extension = model.find_last_of('.'); + if (extension != std::string::npos && extension > 0) + model.erase(extension); + if (!model.empty()) + field("Model", model); + if (models.contains("vae")) + field("VAE", models.at("vae")); + } + if (sampling.contains("eta")) + field("Eta", sampling.at("eta")); + if (info.value("clip_skip", 0) > 0) + field("Clip skip", info.at("clip_skip")); + if (info.contains("strength")) + field("Denoising strength", info.at("strength")); + if (info.contains("hires") && info.at("hires").value("enabled", false)) { + const auto& hires = info.at("hires"); + field("Hires resize", std::to_string(hires.at("target_width").get()) + "x" + std::to_string(hires.at("target_height").get())); + if (hires.value("steps", 0) > 0) + field("Hires steps", hires.at("steps")); + field("Denoising strength", hires.at("denoising_strength")); + field("Hires upscaler", hires.at("upscaler")); + } + field("Version", "stable-diffusion.cpp"); + return result; +} + +std::string generation_info_to_image_metadata(const std::string& info) { + if (info.empty()) + return ""; + const auto parameters = generation_info_to_webui(info); + return parameters.empty() ? info : parameters + ", SDCPP: " + info; +} diff --git a/examples/common/generation_metadata.h b/examples/common/generation_metadata.h new file mode 100644 index 0000000000..7c9e9ffd45 --- /dev/null +++ b/examples/common/generation_metadata.h @@ -0,0 +1,13 @@ +#ifndef __EXAMPLES_COMMON_GENERATION_METADATA_H__ +#define __EXAMPLES_COMMON_GENERATION_METADATA_H__ + +#include +#include + +std::string merge_generation_info(const std::string& source, const std::string& operation); +std::string generation_info_with_relative_lora_paths(const std::string& info, const std::string& lora_model_dir); +int64_t generation_seed(const std::string& info, int64_t fallback = -1); +std::string generation_info_to_webui(const std::string& text); +std::string generation_info_to_image_metadata(const std::string& info); + +#endif // __EXAMPLES_COMMON_GENERATION_METADATA_H__ diff --git a/examples/common/media_io.cpp b/examples/common/media_io.cpp index baac6d2009..e878c889e6 100644 --- a/examples/common/media_io.cpp +++ b/examples/common/media_io.cpp @@ -1,5 +1,7 @@ #include "media_io.h" +#include "generation_metadata.h" #include "log.h" +#include "media_metadata.h" #include "resource_owners.hpp" #include @@ -316,23 +318,22 @@ uint8_t* decode_webp_image_to_buffer(const uint8_t* data, return output.release(); } -std::string build_webp_xmp_packet(const std::string& parameters) { - if (parameters.empty()) { - return ""; - } - - const std::string escaped_parameters = xml_escape(parameters); - return "\n" - "\n" - " \n" - " \n" - " " + - escaped_parameters + - "\n" - " \n" - " \n" - "\n" - ""; +static bool embed_webp_metadata(std::vector& image, const std::string& parameters) { + if (parameters.empty()) + return true; + WebPData data{image.data(), image.size()}; + WebPMuxPtr mux(WebPMuxCreate(&data, 1)); + if (!mux) + return false; + const auto xmp = generation_xmp(parameters); + WebPData xmp_data{reinterpret_cast(xmp.data()), xmp.size()}; + if (WebPMuxSetChunk(mux.get(), "XMP ", &xmp_data, 1) != WEBP_MUX_OK) + return false; + WebPDataGuard assembled; + if (WebPMuxAssemble(mux.get(), &assembled.data) != WEBP_MUX_OK) + return false; + image.assign(assembled.data.bytes, assembled.data.bytes + assembled.data.size); + return true; } bool encode_webp_image_to_vector(const uint8_t* image, @@ -376,37 +377,7 @@ bool encode_webp_image_to_vector(const uint8_t* image, out.assign(encoded.get(), encoded.get() + encoded_size); - if (parameters.empty()) { - return true; - } - - WebPData image_data; - WebPDataInit(&image_data); - WebPDataGuard assembled_data; - - image_data.bytes = out.data(); - image_data.size = out.size(); - - WebPMuxPtr mux(WebPMuxNew()); - if (mux == nullptr) { - return false; - } - - const std::string xmp_packet = build_webp_xmp_packet(parameters); - WebPData xmp_data; - WebPDataInit(&xmp_data); - xmp_data.bytes = reinterpret_cast(xmp_packet.data()); - xmp_data.size = xmp_packet.size(); - - const bool ok = WebPMuxSetImage(mux.get(), &image_data, 1) == WEBP_MUX_OK && - WebPMuxSetChunk(mux.get(), "XMP ", &xmp_data, 1) == WEBP_MUX_OK && - WebPMuxAssemble(mux.get(), &assembled_data.data) == WEBP_MUX_OK; - - if (ok) { - out.assign(assembled_data.data.bytes, assembled_data.data.bytes + assembled_data.data.size); - } - - return ok; + return embed_webp_metadata(out, generation_info_to_image_metadata(parameters)); } #ifdef SD_USE_WEBM @@ -739,6 +710,7 @@ std::vector encode_image_to_vector(EncodedImageFormat format, const std::string& parameters, int quality) { std::vector buffer; + const auto metadata = format == EncodedImageFormat::WEBP ? std::string() : generation_info_to_image_metadata(parameters); auto write_func = [&buffer](void* context, void* data, int size) { (void)context; @@ -761,7 +733,7 @@ std::vector encode_image_to_vector(EncodedImageFormat format, result = stbi_write_jpg_to_func(c_func, &ctx, width, height, channels, image, quality); break; case EncodedImageFormat::PNG: - result = stbi_ext_write_png_to_func(c_func, &ctx, width, height, channels, image, width * channels, parameters.empty() ? nullptr : parameters.c_str()); + result = stbi_ext_write_png_to_func(c_func, &ctx, width, height, channels, image, width * channels, metadata.empty() ? nullptr : metadata.c_str()); break; case EncodedImageFormat::WEBP: #ifdef SD_USE_WEBP @@ -779,9 +751,12 @@ std::vector encode_image_to_vector(EncodedImageFormat format, break; } - if (!result) { - buffer.clear(); + if (result && format == EncodedImageFormat::JPEG && !embed_jpeg_metadata(buffer, metadata)) { + LOG_ERROR("generation metadata exceeds the container capacity or the encoded image is invalid"); + result = 0; } + if (!result) + buffer.clear(); return buffer; } @@ -794,18 +769,8 @@ bool write_image_to_file(const std::string& path, int quality) { const EncodedImageFormat format = encoded_image_format_from_path(path); - switch (format) { - case EncodedImageFormat::JPEG: - return stbi_write_jpg(path.c_str(), width, height, channels, image, quality, parameters.empty() ? nullptr : parameters.c_str()) != 0; - case EncodedImageFormat::PNG: - return stbi_write_png(path.c_str(), width, height, channels, image, 0, parameters.empty() ? nullptr : parameters.c_str()) != 0; - case EncodedImageFormat::WEBP: { - const std::vector encoded = encode_image_to_vector(format, image, width, height, channels, parameters, quality); - return !encoded.empty() && write_binary_file_bytes(path, encoded); - } - default: - return false; - } + const auto encoded = encode_image_to_vector(format, image, width, height, channels, parameters, quality); + return !encoded.empty() && write_binary_file_bytes(path, encoded); } uint8_t* load_image_from_file(const char* image_path, @@ -1384,7 +1349,10 @@ std::vector create_video_from_sd_images_to_vector(const std::string& ou #ifdef SD_USE_WEBP if (format == "webp") { - return create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality); + auto data = create_animated_webp_from_sd_images_to_vector(images, num_images, fps, quality); + if (!data.empty() && !embed_webp_metadata(data, parameters)) + return {}; + return data; } #endif @@ -1469,70 +1437,94 @@ static uint32_t read_le32(const uint8_t* data) { (static_cast(data[3]) << 24); } +/** Loads a WAV file from disk and delegates decoding to load_wav_from_memory(). */ bool load_wav_from_file(const std::string& path, std::vector& interleaved_samples, uint32_t& sample_rate, uint32_t& channels) { + std::ifstream file(fs::u8path(path), std::ios::binary); + if (!file) { + interleaved_samples.clear(); + sample_rate = 0; + channels = 0; + return false; + } + file.seekg(0, std::ios::end); + const std::streamoff file_size = file.tellg(); + file.seekg(0, std::ios::beg); + if (file_size < 0) { + interleaved_samples.clear(); + sample_rate = 0; + channels = 0; + return false; + } + std::vector buffer(static_cast(file_size)); + if (file_size > 0 && !file.read(reinterpret_cast(buffer.data()), file_size)) { + interleaved_samples.clear(); + sample_rate = 0; + channels = 0; + return false; + } + return load_wav_from_memory(buffer.data(), buffer.size(), interleaved_samples, sample_rate, channels); +} + +/** Parses a RIFF/WAVE buffer in place; compares chunk sizes against the remaining + * bytes so payload-controlled sizes cannot wrap size_t on 32-bit builds. */ +bool load_wav_from_memory(const uint8_t* data, + size_t size, + std::vector& interleaved_samples, + uint32_t& sample_rate, + uint32_t& channels) { interleaved_samples.clear(); sample_rate = 0; channels = 0; - std::ifstream file(fs::u8path(path), std::ios::binary); - uint8_t riff_header[12]; - if (!file.read(reinterpret_cast(riff_header), sizeof(riff_header)) || - std::memcmp(riff_header, "RIFF", 4) != 0 || - std::memcmp(riff_header + 8, "WAVE", 4) != 0) { + if (data == nullptr || + size < 12 || + std::memcmp(data, "RIFF", 4) != 0 || + std::memcmp(data + 8, "WAVE", 4) != 0) { return false; } uint16_t audio_format = 0; uint16_t bits_per_sample = 0; uint16_t block_align = 0; - std::streampos data_pos = std::streampos(-1); + size_t data_pos = SIZE_MAX; uint32_t data_size = 0; - while (file.good()) { - uint8_t chunk_header[8]; - if (!file.read(reinterpret_cast(chunk_header), sizeof(chunk_header))) { - break; - } - uint32_t chunk_size = read_le32(chunk_header + 4); - std::streampos chunk_data_pos = file.tellg(); - - if (std::memcmp(chunk_header, "fmt ", 4) == 0) { - if (chunk_size < 16) { - return false; - } - std::vector fmt(chunk_size); - if (!file.read(reinterpret_cast(fmt.data()), chunk_size)) { + size_t pos = 12; + while (pos + 8 <= size) { + uint32_t chunk_size = read_le32(data + pos + 4); + size_t chunk_data_pos = pos + 8; + const size_t remaining = size - chunk_data_pos; // safe: chunk_data_pos <= size + // Chunk sizes come from the payload, so compare against the remaining bytes + // instead of adding them to offsets (avoids size_t wrap on 32-bit builds). + const bool data_fits = static_cast(chunk_size) <= remaining; + + if (std::memcmp(data + pos, "fmt ", 4) == 0) { + if (chunk_size < 16 || !data_fits) { return false; } - audio_format = read_le16(fmt.data()); - channels = read_le16(fmt.data() + 2); - sample_rate = read_le32(fmt.data() + 4); - block_align = read_le16(fmt.data() + 12); - bits_per_sample = read_le16(fmt.data() + 14); + audio_format = read_le16(data + chunk_data_pos); + channels = read_le16(data + chunk_data_pos + 2); + sample_rate = read_le32(data + chunk_data_pos + 4); + block_align = read_le16(data + chunk_data_pos + 12); + bits_per_sample = read_le16(data + chunk_data_pos + 14); if (audio_format == 0xfffe && chunk_size >= 40) { - audio_format = read_le16(fmt.data() + 24); + audio_format = read_le16(data + chunk_data_pos + 24); } - } else if (std::memcmp(chunk_header, "data", 4) == 0) { + } else if (std::memcmp(data + pos, "data", 4) == 0) { data_pos = chunk_data_pos; data_size = chunk_size; - file.seekg(chunk_size, std::ios::cur); - } else { - file.seekg(chunk_size, std::ios::cur); - } - - if (!file.good()) { - break; - } - if ((chunk_size & 1) != 0) { - file.seekg(1, std::ios::cur); } + // Only advance past a chunk whose size and pad byte fully fit (keeps the sum + // overflow-free); a chunk truncated at the buffer end terminates the walk. + const bool chunk_fits = data_fits && ((chunk_size & 1) == 0 || static_cast(chunk_size) < remaining); + pos = chunk_fits ? chunk_data_pos + static_cast(chunk_size) + (chunk_size & 1) : size; } const uint32_t bytes_per_sample = (bits_per_sample + 7) / 8; - if (data_pos == std::streampos(-1) || data_size == 0 || channels == 0 || sample_rate == 0 || + if (data_pos == SIZE_MAX || data_size == 0 || channels == 0 || sample_rate == 0 || block_align == 0 || bytes_per_sample == 0 || block_align < channels * bytes_per_sample || (audio_format != 1 && audio_format != 3)) { return false; @@ -1542,16 +1534,14 @@ bool load_wav_from_file(const std::string& path, if (frame_count == 0 || frame_count > SIZE_MAX / channels) { return false; } - std::vector pcm(data_size); - file.clear(); - file.seekg(data_pos); - if (!file.read(reinterpret_cast(pcm.data()), data_size)) { + if (static_cast(data_size) > size - data_pos) { return false; } + const uint8_t* pcm = data + data_pos; interleaved_samples.resize(static_cast(frame_count * channels)); for (uint64_t frame = 0; frame < frame_count; ++frame) { - const uint8_t* frame_data = pcm.data() + frame * block_align; + const uint8_t* frame_data = pcm + frame * block_align; for (uint32_t channel = 0; channel < channels; ++channel) { const uint8_t* sample_data = frame_data + channel * bytes_per_sample; float sample = 0.0f; diff --git a/examples/common/media_io.h b/examples/common/media_io.h index b2b56eccde..7feb0ada25 100644 --- a/examples/common/media_io.h +++ b/examples/common/media_io.h @@ -16,6 +16,7 @@ enum class EncodedImageFormat { EncodedImageFormat encoded_image_format_from_path(const std::string& path); +// parameters is the core sdcpp.generation/v1 JSON; empty disables both metadata formats. std::vector encode_image_to_vector(EncodedImageFormat format, const uint8_t* image, int width, @@ -125,4 +126,11 @@ bool load_wav_from_file(const std::string& path, uint32_t& sample_rate, uint32_t& channels); +/** Decodes a WAV buffer (PCM 8/16/24/32 or IEEE float 32/64) into interleaved samples. */ +bool load_wav_from_memory(const uint8_t* data, + size_t size, + std::vector& interleaved_samples, + uint32_t& sample_rate, + uint32_t& channels); + #endif // __MEDIA_IO_H__ diff --git a/examples/common/media_metadata.cpp b/examples/common/media_metadata.cpp new file mode 100644 index 0000000000..d5e9af3fa3 --- /dev/null +++ b/examples/common/media_metadata.cpp @@ -0,0 +1,47 @@ +#include "media_metadata.h" + +std::string generation_xmp(const std::string& info) { + std::string escaped; + for (char c : info) { + switch (c) { + case '&': + escaped += "&"; + break; + case '<': + escaped += "<"; + break; + case '>': + escaped += ">"; + break; + case '"': + escaped += """; + break; + case '\'': + escaped += "'"; + break; + default: + escaped += c; + break; + } + } + return "" + "" + "" + + escaped + ""; +} + +bool embed_jpeg_metadata(std::vector& image, const std::string& parameters) { + if (parameters.empty()) + return true; + const char label[] = "parameters"; + const size_t overhead = sizeof(label) + 1; + if (image.size() < 2 || parameters.size() > 65533 - overhead) + return false; + const size_t length = overhead + parameters.size() + 2; + std::vector segment = {0xff, 0xfe, static_cast(length >> 8), static_cast(length)}; + segment.insert(segment.end(), label, label + sizeof(label)); + segment.insert(segment.end(), parameters.begin(), parameters.end()); + segment.push_back(0); + image.insert(image.begin() + 2, segment.begin(), segment.end()); + return true; +} diff --git a/examples/common/media_metadata.h b/examples/common/media_metadata.h new file mode 100644 index 0000000000..1843dee19d --- /dev/null +++ b/examples/common/media_metadata.h @@ -0,0 +1,11 @@ +#ifndef __EXAMPLES_COMMON_MEDIA_METADATA_H__ +#define __EXAMPLES_COMMON_MEDIA_METADATA_H__ + +#include +#include +#include + +std::string generation_xmp(const std::string& parameters); +bool embed_jpeg_metadata(std::vector& image, const std::string& parameters); + +#endif // __EXAMPLES_COMMON_MEDIA_METADATA_H__ diff --git a/examples/common/resource_owners.hpp b/examples/common/resource_owners.hpp index 73f8bdd81d..8a8f843146 100644 --- a/examples/common/resource_owners.hpp +++ b/examples/common/resource_owners.hpp @@ -8,6 +8,7 @@ #include #include +#include "generation_metadata.h" #include "stable-diffusion.h" struct FreeDeleter { @@ -175,6 +176,7 @@ class SDAudioOwner { class SDImageVec { private: std::vector images_; + std::vector generation_infos_; public: SDImageVec() = default; @@ -183,13 +185,14 @@ class SDImageVec { SDImageVec& operator=(const SDImageVec&) = delete; SDImageVec(SDImageVec&& other) noexcept - : images_(std::move(other.images_)) { + : images_(std::move(other.images_)), generation_infos_(std::move(other.generation_infos_)) { } SDImageVec& operator=(SDImageVec&& other) noexcept { if (this != &other) { clear(); - images_ = std::move(other.images_); + images_ = std::move(other.images_); + generation_infos_ = std::move(other.generation_infos_); } return *this; } @@ -210,6 +213,41 @@ class SDImageVec { images_.reserve(count); } + const std::string& generation_info(size_t index) const { + static const std::string empty; + return index < generation_infos_.size() ? generation_infos_[index] : empty; + } + + void append_generation_info(size_t index, const std::string& info, const std::string& lora_model_dir = {}) { + generation_infos_.resize(images_.size()); + generation_infos_.at(index) = merge_generation_info(generation_infos_.at(index), generation_info_with_relative_lora_paths(info, lora_model_dir)); + } + + void adopt(sd_image_generation_result_t* result, const std::string& lora_model_dir = {}) { + clear(); + if (!result) + return; + for (int i = 0; i < result->count; ++i) { + images_.push_back(result->items[i].image); + result->items[i].image = {}; + generation_infos_.push_back(generation_info_with_relative_lora_paths(result->items[i].generation_info, lora_model_dir)); + } + free_sd_image_generation_result(result); + } + + void adopt_video(sd_video_generation_result_t* result, const std::string& lora_model_dir = {}) { + clear(); + if (!result) + return; + const auto info = generation_info_with_relative_lora_paths(result->generation_info, lora_model_dir); + for (int i = 0; i < result->frame_count; ++i) { + images_.push_back(result->frames[i]); + result->frames[i] = {}; + generation_infos_.push_back(info); + } + free_sd_video_generation_result(result); + } + void adopt(sd_image_t* images, int count) { clear(); if (images == nullptr || count <= 0) { @@ -270,6 +308,7 @@ class SDImageVec { image.data = nullptr; } images_.clear(); + generation_infos_.clear(); } }; diff --git a/examples/server/CMakeLists.txt b/examples/server/CMakeLists.txt index 58260be0d4..d01f7a7e20 100644 --- a/examples/server/CMakeLists.txt +++ b/examples/server/CMakeLists.txt @@ -64,8 +64,10 @@ endif() add_executable(${TARGET} ../common/common.cpp + ../common/generation_metadata.cpp ../common/log.cpp ../common/media_io.cpp + ../common/media_metadata.cpp main.cpp runtime.cpp model_catalog.cpp diff --git a/examples/server/README.md b/examples/server/README.md index 1bb49a59f5..8d8f687cb5 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -89,8 +89,8 @@ FLUX 1/2/3, Chroma, Qwen Image and editing variants; Z-Image, Anima, ERNIE, LENS, PiD and other supported image families; AnimateDiff, Wan, LTX, HunyuanVideo, LingBot Video and MiniMax-H3. Model-specific text encoder labels identify the expected encoder family and size. Loading a video model enables **Video Generation**; I2V also requires a -start image. The Models panel can be collapsed and remembers its open state. -Its header shows elapsed loading time from the status API and keeps the final +start image. The model button below the prompt opens model and recipe selection. +Its status summary shows elapsed loading time from the status API and keeps the final time after success or failure, including after a page refresh. Built-ins are embedded in the binary from [recipes/](recipes/), with one JSON file per recipe. Adding or diff --git a/examples/server/api.md b/examples/server/api.md index bb42523661..97cbd66cdf 100644 --- a/examples/server/api.md +++ b/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: @@ -183,6 +183,9 @@ Default-enabled features are `img2img`, `mask_image`, `lora`, `vae_tiling`, | `clip_skip` | CLIP Skip | | `img2img`, `mask_image` | Image initialization/strength and masks | | `ref_images` | Reference image inputs | +| `ref_videos` | Reference videos as frame arrays | +| `ref_audios` | Standalone reference audio; runtime requires an Audio VAE or Audio Encoder | +| `ref_video_audios` | Soundtracks paired with reference videos; runtime requires an Audio VAE | | `controlnet` | Control image and ControlNet strength | | `control_image` | Control image without a separate ControlNet (FLUX Canny/Depth) | | `ip_adapter` | IP-Adapter image and strength | @@ -241,6 +244,7 @@ the same one-recipe-per-file schema: "id": "my-flux", "name": "My FLUX", "description": "FLUX.1 dev", + "size_multiple": 16, "model_url": "https://huggingface.co/black-forest-labs/FLUX.1-dev", "slots": [ {"id": "diffusion_model", "required": true}, @@ -256,7 +260,8 @@ the same one-recipe-per-file schema: IDs contain 1–128 ASCII letters, digits, dots, underscores or hyphens. `name`, `schema_version: 1` and `slots` are required; `description` and -`model_url`, `required_any`, `defaults`, `context`, `features` and `disabled_features` are optional. Each slot requires only `id` and boolean `required`. +`model_url`, `required_any`, `defaults`, `context`, `features`, `disabled_features` +and `size_multiple`, `video_constraints` are optional. Each slot requires only `id` and boolean `required`. The server supplies the default label and search directories for each slot. An optional nonempty `label` overrides the display name, for example `{"id": "diffusion_model", "label": "Low-noise diffusion model", "required": true}`. @@ -264,6 +269,53 @@ 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 +generated size presets multiples of 16. Exact dimensions that do not match the +increment require correction before submission. This is recipe metadata, not a +generation parameter, and does not change the core model's alignment behavior. + +`video_constraints` is optional metadata for the web UI's Frames / Duration +controls. It does not override the runtime's model-specific rules. For example: + +```json +"video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 +} +``` + +Legal frame counts are `frame_step * n + frame_offset`, with integer `n >= 0`, +within the inclusive `min_frames` / `max_frames` range. All fields are optional: + +| Field | Default | Accepted values | +| --- | --- | --- | +| `frame_step` | `1` | Integer 1–4096 | +| `frame_offset` | `0` | Integer 0 to `frame_step - 1` | +| `frame_rounding` | `"floor"` | `"floor"` or `"ceil"` | +| `min_frames` | `1` | Integer 1–4096 | +| `max_frames` | `4096` | Integer 1–4096, at least `min_frames` | +| `fixed_fps` | Unset | Integer 1–240; locks the FPS control | + +The range must contain at least one legal frame count. Unknown keys are rejected. +Omitting the object adds no frame alignment in the UI. `defaults.fps` only sets +the initial FPS; it does not lock it. The UI multiplies requested duration by FPS, +aligns using the declared rounding direction, then clamps to the legal range. +Duration stays in the UI; generation requests still send only `video_frames` and +`fps`. Switching input modes preserves the aligned length. Changing FPS keeps +frames in Frames mode, and keeps the requested seconds in Duration mode. + `model_url` is an optional HTTPS Hugging Face model repository link, shown beside the recipe description. Built-ins link to the model's author/publisher, not repacks. `required_any` contains groups of declared slot IDs. At least one model in each @@ -280,7 +332,8 @@ TAE slot. | `model` | `checkpoints` | | `diffusion_model`, `high_noise_diffusion_model`, `uncond_diffusion_model` | `diffusion_models` | | `clip_l`, `clip_g`, `t5xxl`, `llm`, `audio_encoder`, `embeddings_connectors` | `text_encoders` | -| `clip_vision`, `llm_vision` | `clip_vision` | +| `clip_vision` | `clip_vision` | +| `llm_vision` | `text_encoders` or `clip_vision` | | `vae`, `audio_vae` | `vae` | | `tae` (legacy alias: `taesd`) | `vae` or `vae_approx` | | `tokenizer` | `tokenizers` | @@ -302,11 +355,15 @@ new model architectures. Optional `context` settings apply to each model load and override the matching startup setting. Supported fields are `vae_format` (`auto`, `flux`, `sd3`, -`flux2`, or `wan`) and `model_args` (a string in the same format as -`--model-args`; replaces the startup string). Other startup settings remain -unchanged. Switching to a recipe that omits these fields restores the startup -values. For example, PiD uses `{"vae_format": "flux"}` and Qwen Image Edit -2511 uses `{"model_args": "qwen_image_zero_cond_t=true"}`. +`flux2`, or `wan`), `model_args` (a string in the same format as +`--model-args`; replaces the startup string), and `backend` / `params_backend` +(strings in the same format as `--backend` / `--params-backend`; replace the +startup assignments for that load, so a recipe can place components on +specific devices, e.g. keep a large text encoder on CPU while another recipe +uses the GPU). Other startup settings remain unchanged. Switching to a +recipe that omits these fields restores the startup values. For example, PiD +uses `{"vae_format": "flux"}` and Qwen Image Edit 2511 uses +`{"model_args": "qwen_image_zero_cond_t=true"}`. Supported defaults: @@ -442,10 +499,6 @@ image establishes the generation dimensions if `size` is omitted. Input geometry follows `image_preprocess`: references preserve their dimensions by default, while init and mask use the generation canvas preset. -When the request starts without preconfigured reference images, the first uploaded image becomes the img2img init image if the loaded model does not support reference-image conditioning, the final reference list is empty, a mask is present, or an explicit numeric `strength` override is provided. Otherwise, only the final reference list is used for image conditioning, without initializing the img2img latent. Requests with preconfigured references retain their existing init-image selection. - -Routing uses the reference list and mask after applying `sd_cpp_extra_args`. Both `ref_images: []` and `ref_images: null` clear the reference list; `mask_image: null` clears an uploaded mask. An explicit `init_image` in `sd_cpp_extra_args`, including `null`, takes precedence over the implicit init-image selection. - Reference encoding then follows model presets and `ref_image_args`. To skip input geometry for references and disable resizing before VAE encoding, include this in `prompt`: @@ -872,7 +925,7 @@ Shared default fields used by both `img_gen` and `vid_gen`: `vae_tiling_params.tile_size_w` and `tile_size_h` are in **image pixels**, with `0` selecting the 256-pixel default. Both encode and decode use these sizes without an encoding multiplier. Positive `rel_size_w`/`rel_size_h` values override the corresponding absolute size: values up to 1 are dimension fractions, and values greater than 1 are target tile counts. Set `enabled` to use spatial tiling. Sizes are aligned down to the VAE scale factor and capped at the input dimensions; explicit sizes below the minimum supported tile size are rejected. These fields previously used latent units; see [VAE tiling](../../docs/performance.md#use-vae-tiling-to-reduce-encode-and-decode-memory-usage) for migration and OOM retry behavior. -`vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`). +`vae_tiling_params.extra_tiling_args` accepts a key=value list. Supported video VAEs accept `temporal_tile_frames` (alias `temporal_tile_size`, default `4`) and `temporal_tile_overlap` (default `1`). Wan2.2 TI2V 5B uses `temporal_decode_frames` (default `1`) to decode multiple latent frames per forward pass, trading additional memory for throughput; other model versions ignore this option. LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEHV use overlap blending. MiniMax H3 keeps its model-specific fixed temporal windows because its latent-to-frame mapping is non-linear. `img_gen`-specific default fields: @@ -901,6 +954,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH | --- | --- | | `video_frames` | `integer` | | `fps` | `integer` | +| `ref_image_args` | `string` | | `moe_boundary` | `number` | | `vace_strength` | `number` | | `high_noise_sample_params` | `object` | @@ -1032,7 +1086,11 @@ Notes: #### `POST /sdcpp/v1/jobs/{id}/cancel` -Attempts to cancel an accepted job. +Cancels a queued job immediately, or requests interruption of a generating image or +video job. A generating job returns `200 OK` with `cancel_requested: true` and +remains `generating` until execution stops at a cancellation checkpoint. Continue +polling until its status becomes `cancelled`. Cancelled jobs have no result. +Repeated cancellation requests are safe; terminal jobs are returned unchanged. Typical status codes: @@ -1394,7 +1452,7 @@ Compared with `img_gen`, the `vid_gen` request body: - `vid_gen` is a single video sequence job, so `batch_count` is not part of the request schema - `mask_image`, `control_image`, `control_strength`, `ip_adapter_image`, `ip_adapter_strength`, and `embed_image_metadata` are not part of the request schema - `ref_images` is accepted for MiniMax-H3 Ref2VA conditioning; other video model families currently ignore it -- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `moe_boundary`, and `vace_strength` +- `vid_gen` adds `end_image`, `control_frames`, `high_noise_sample_params`, `video_frames`, `fps`, `ref_image_args`, `moe_boundary`, and `vace_strength` Example: @@ -1409,6 +1467,7 @@ Example: "seed": -1, "video_frames": 33, "fps": 16, + "ref_image_args": "", "moe_boundary": 0.875, "vace_strength": 1.0, @@ -1490,6 +1549,7 @@ Example: - Images retain request order and correspond to ``, ``, and so on in the prompt. - For MiniMax-H3, non-empty `ref_images` cannot be combined with `init_image` or `end_image`. A conflicting API request fails during generation; the WebUI checks this before submission. - MiniMax-H3 does not support `control_frames`; leave that array empty. +- `ref_image_args` applies to `vid_gen`; MiniMax-H3 honors `ref_image_size` there (`match` default, `max` for a 2048 px short edge). Only references above the limit are downscaled; each axis is then snapped to the nearest 32 px multiple. - Other video model families currently ignore `ref_images`. - `features_by_mode.vid_gen` does not currently advertise `ref_images`; its absence is not an indication that MiniMax-H3 reference images are unsupported. @@ -1540,6 +1600,7 @@ Top-level scalar fields: | `seed` | `integer` | | `video_frames` | `integer` | | `fps` | `integer` | +| `ref_image_args` | `string` | | `moe_boundary` | `number` | | `vace_strength` | `number` | | `preview` | `string` | @@ -1681,15 +1742,26 @@ When omitted, backend defaults apply to these fields: ### Frame Count Semantics -`video_frames` is the requested target length, but the current core video path internally normalizes the effective frame count to the largest `4n + 1` value that does not exceed the requested count. +`video_frames` is the requested target length. The core normalizes it according +to the loaded model: -Examples: +| Model family | Frame alignment | Fixed FPS | +| --- | --- | --- | +| Wan, HunyuanVideo, LingBot | `4n + 1`, round down | Wan 2.2 S2V: 16 | +| LTX | `8n + 1`, round down | None | +| MiniMax-H3 | `17n + 5`, round up, minimum 5 | 24 | +| AnimateDiff | Whole frames, maximum 32 | None | + +Wan examples: - `video_frames = 33` stays `33` - `video_frames = 34` becomes `33` - `video_frames = 32` becomes `29` -The completed job payload includes the actual decoded `frame_count`. +At 16 FPS, a 5-second target requests 80 frames and aligns to 77 frames for Wan, +or approximately 4.813 seconds. Audio-driven generation and temporal upscaling +can change the final length. Use the completed job's actual decoded `frame_count` +and `fps` for the final playback duration. ### Completion Result @@ -1767,3 +1839,14 @@ Example cancelled job: - `400 Bad Request` for an empty body, unsupported model mode, invalid JSON, invalid generation parameters, or an unsupported output format - `429 Too Many Requests` when the job queue is full - `500 Internal Server Error` for unexpected server exceptions during submission + +## Actual generation parameters + +Generation responses expose the core `sdcpp.generation/v1` record as +`generation_info`, independently of `embed_image_metadata`. Native async image +result items and OpenAI image items contain the parameter snapshot with each +image's actual seed and without `batch_count`. Native video +results contain the video request snapshot. SD API `info.generation_info` is an array +aligned with `info.infotexts` and `info.all_seeds`; the outer `parameters` field +remains the request echo. See [generation metadata](../../docs/generation_metadata.md) +for field semantics, postprocessing and container storage. diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 7f3c8a1bfc..7a80107897 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -3,13 +3,46 @@ #include "async_jobs.h" #include +#include #include #include +#include #include "common/log.h" #include "common/media_io.h" #include "common/resource_owners.hpp" +class JobCancellationMonitor { + std::mutex mutex_; + std::condition_variable cv_; + bool done_ = false; + std::thread thread_; + +public: + // The caller holds sd_ctx_mutex until this monitor has joined, keeping the + // context alive and preventing cancellation from reaching another request. + JobCancellationMonitor(sd_ctx_t* ctx, const AsyncGenerationJob& job) + : thread_([this, ctx, &job]() { + std::unique_lock lock(mutex_); + while (!done_) { + if (job.cancel_requested.load()) { + // Each generation pass resets the library flag, including ADetailer. + sd_cancel_generation(ctx, SD_CANCEL_ALL); + } + cv_.wait_for(lock, std::chrono::milliseconds(10), [this]() { return done_; }); + } + }) {} + + ~JobCancellationMonitor() { + { + std::lock_guard lock(mutex_); + done_ = true; + } + cv_.notify_one(); + thread_.join(); + } +}; + static void preview_callback(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data) { (void)is_noisy; if (frame_count <= 0 || frames == nullptr || frames[0].data == nullptr) { @@ -142,9 +175,11 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { } manager.queue.erase(new_end, manager.queue.end()); + job.cancel_requested.store(true); job.status = AsyncJobStatus::Cancelled; job.completed_at = unix_timestamp_now_ms(); job.result_images_b64.clear(); + job.result_generation_info = nullptr; job.result_media_b64.clear(); job.result_media_mime_type.clear(); job.result_frame_count = 0; @@ -156,13 +191,14 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJob& job) { json result; - result["id"] = job.id; - result["kind"] = async_job_kind_name(job.kind); - result["status"] = async_job_status_name(job.status); - result["created"] = job.created_at; - result["started"] = job.started_at == 0 ? json(nullptr) : json(job.started_at); - result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); - result["queue_position"] = 0; + result["id"] = job.id; + result["kind"] = async_job_kind_name(job.kind); + result["status"] = async_job_status_name(job.status); + result["cancel_requested"] = job.cancel_requested.load(); + result["created"] = job.created_at; + result["started"] = job.started_at == 0 ? json(nullptr) : json(job.started_at); + result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); + result["queue_position"] = 0; if (job.status == AsyncJobStatus::Queued) { size_t position = 1; @@ -183,11 +219,12 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo {"fps", job.result_fps}, {"frame_count", job.result_frame_count}, {"b64_json", job.result_media_b64}, + {"generation_info", job.result_generation_info}, }; } else { json images = json::array(); for (size_t i = 0; i < job.result_images_b64.size(); ++i) { - images.push_back({{"index", i}, {"b64_json", job.result_images_b64[i]}}); + images.push_back({{"index", i}, {"b64_json", job.result_images_b64[i]}, {"generation_info", job.result_generation_info.at(i)}}); } result["result"] = { {"output_format", job.img_gen.output_format}, @@ -199,11 +236,11 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo job.status == AsyncJobStatus::Cancelled) { result["result"] = nullptr; result["error"] = { - {"code", + {"code", job.error_code.empty() - ? (job.status == AsyncJobStatus::Cancelled ? "cancelled" : "generation_failed") - : job.error_code}, - {"message", job.error_message}, + ? (job.status == AsyncJobStatus::Cancelled ? "cancelled" : "generation_failed") + : job.error_code}, + {"message", job.error_message}, }; } else { result["result"] = nullptr; @@ -228,26 +265,35 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo bool execute_img_gen_job(ServerRuntime& runtime, AsyncGenerationJob& job, std::vector& output_images, + json& output_info, std::string& error_message) { + output_info = json::array(); sd_img_gen_params_t params = job.img_gen.to_sd_img_gen_params_t(); SDImageVec results; { std::lock_guard lock(*runtime.sd_ctx_mutex); - set_preview_callback(job); - sd_image_t* raw_results = nullptr; - int num_results = 0; - if (!generate_image(runtime.sd_ctx, ¶ms, &raw_results, &num_results)) { - raw_results = nullptr; - num_results = 0; + if (job.cancel_requested.load()) { + return false; } + JobCancellationMonitor cancellation(runtime.sd_ctx, job); + set_preview_callback(job); + sd_image_generation_result_t* generated = nullptr; + generate_image(runtime.sd_ctx, ¶ms, &generated); + results.adopt(generated, runtime.ctx_params->lora_model_dir); clear_preview_callback(); - results.adopt(raw_results, num_results); + 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; + } } const int num_results = results.count(); - if (num_results <= 0) { + if (job.cancel_requested.load() || num_results <= 0) { error_message = "generate_image returned no results"; return false; } @@ -259,17 +305,13 @@ bool execute_img_gen_job(ServerRuntime& runtime, encoded_format = EncodedImageFormat::WEBP; } - int batch_count = job.img_gen.gen_params.batch_count; - int images_per_batch = batch_count > 0 ? std::max(1, num_results / batch_count) : 1; for (int i = 0; i < num_results; ++i) { if (results[i].data == nullptr) { continue; } const std::string metadata = job.img_gen.gen_params.embed_image_metadata - ? get_image_params(*runtime.ctx_params, - job.img_gen.gen_params, - job.img_gen.gen_params.seed + i / images_per_batch) + ? results.generation_info(i) : ""; auto image_bytes = encode_image_to_vector(encoded_format, results[i].data, @@ -282,6 +324,7 @@ bool execute_img_gen_job(ServerRuntime& runtime, continue; } output_images.push_back(base64_encode(image_bytes)); + output_info.push_back(json::parse(results.generation_info(i))); } if (output_images.empty()) { @@ -298,11 +341,10 @@ bool execute_vid_gen_job(ServerRuntime& runtime, std::string& output_media_mime_type, int& output_frame_count, int& output_fps, + json& output_info, std::string& error_message) { sd_vid_gen_params_t params = job.vid_gen.to_sd_vid_gen_params_t(); - std::string str_params = job.vid_gen.gen_params.embed_image_metadata - ? get_image_params(*runtime.ctx_params, job.vid_gen.gen_params, job.vid_gen.gen_params.seed, VID_GEN) - : ""; + std::string str_params; SDImageVec results; int num_results = 0; @@ -310,17 +352,26 @@ bool execute_vid_gen_job(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); + if (job.cancel_requested.load()) { + return false; + } + JobCancellationMonitor cancellation(runtime.sd_ctx, job); set_preview_callback(job); - sd_image_t* raw_results = nullptr; - if (!generate_video(runtime.sd_ctx, ¶ms, &raw_results, &num_results, &generated_audio, &output_fps)) { - raw_results = nullptr; + sd_video_generation_result_t* generated = nullptr; + if (generate_video(runtime.sd_ctx, ¶ms, &generated)) { + output_fps = generated->fps; + generated_audio = generated->audio; + generated->audio = nullptr; + results.adopt_video(generated, runtime.ctx_params->lora_model_dir); + output_info = json::parse(results.generation_info(0)); + if (job.vid_gen.gen_params.embed_image_metadata) + str_params = results.generation_info(0); } clear_preview_callback(); - results.adopt(raw_results, num_results); } num_results = results.count(); - if (num_results <= 0) { + if (job.cancel_requested.load() || num_results <= 0) { free_sd_audio(generated_audio); error_message = "generate_video returned no results"; return false; @@ -377,6 +428,7 @@ void async_job_worker(ServerRuntime& runtime) { } std::vector output_images; + json output_info; std::string output_media_b64; std::string output_media_mime_type; int output_frame_count = 0; @@ -385,7 +437,7 @@ void async_job_worker(ServerRuntime& runtime) { bool ok = false; if (job->kind == AsyncJobKind::ImgGen) { - ok = execute_img_gen_job(runtime, *job, output_images, error_message); + ok = execute_img_gen_job(runtime, *job, output_images, output_info, error_message); } else if (job->kind == AsyncJobKind::VidGen) { ok = execute_vid_gen_job(runtime, *job, @@ -393,6 +445,7 @@ void async_job_worker(ServerRuntime& runtime) { output_media_mime_type, output_frame_count, output_fps, + output_info, error_message); } else { error_message = "unsupported job kind"; @@ -406,9 +459,14 @@ void async_job_worker(ServerRuntime& runtime) { } job->completed_at = unix_timestamp_now_ms(); - if (ok) { + if (job->cancel_requested.load()) { + job->status = AsyncJobStatus::Cancelled; + job->error_code = "cancelled"; + job->error_message = "job cancelled by client"; + } else if (ok) { job->status = AsyncJobStatus::Completed; job->result_images_b64 = std::move(output_images); + job->result_generation_info = std::move(output_info); job->result_media_b64 = std::move(output_media_b64); job->result_media_mime_type = std::move(output_media_mime_type); job->result_frame_count = output_frame_count; @@ -420,6 +478,7 @@ void async_job_worker(ServerRuntime& runtime) { job->error_code = "generation_failed"; job->error_message = error_message.empty() ? "unknown generation error" : error_message; job->result_images_b64.clear(); + job->result_generation_info = nullptr; job->result_media_b64.clear(); job->result_media_mime_type.clear(); job->result_frame_count = 0; diff --git a/examples/server/async_jobs.h b/examples/server/async_jobs.h index 049a9d73b3..6fa3cbd138 100644 --- a/examples/server/async_jobs.h +++ b/examples/server/async_jobs.h @@ -1,5 +1,6 @@ #pragma once +#include #include #include #include @@ -32,12 +33,14 @@ struct AsyncGenerationJob { std::string id; AsyncJobKind kind = AsyncJobKind::ImgGen; AsyncJobStatus status = AsyncJobStatus::Queued; - int64_t created_at = unix_timestamp_now_ms(); - int64_t started_at = 0; - int64_t completed_at = 0; + std::atomic cancel_requested{false}; + int64_t created_at = unix_timestamp_now_ms(); + int64_t started_at = 0; + int64_t completed_at = 0; ImgGenJobRequest img_gen; VidGenJobRequest vid_gen; std::vector result_images_b64; + json result_generation_info; std::string result_media_b64; std::string result_media_mime_type; int result_frame_count = 0; @@ -72,6 +75,7 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo bool execute_img_gen_job(ServerRuntime& runtime, AsyncGenerationJob& job, std::vector& output_images, + json& output_info, std::string& error_message); bool execute_vid_gen_job(ServerRuntime& runtime, AsyncGenerationJob& job, @@ -79,5 +83,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime, std::string& output_media_mime_type, int& output_frame_count, int& output_fps, + json& output_info, std::string& error_message); void async_job_worker(ServerRuntime& runtime); diff --git a/examples/server/frontend b/examples/server/frontend index e36e34a2e2..fa93429c47 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit e36e34a2e2204e3df654b16cadc112d8ce8f3174 +Subproject commit fa93429c47c924054e6a9c14f8d7b8e06f57a0ec diff --git a/examples/server/main.cpp b/examples/server/main.cpp index 170877686b..a10cb66f1c 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -48,6 +48,10 @@ static void parse_args(int argc, log_level = svr_params.log_level; log_color = svr_params.color; + if (ctx_params.list_devices) { + exit(ctx_params.print_devices() ? 0 : 1); + } + const bool random_seed_requested = default_gen_params.seed < 0; if (!svr_params.resolve_and_validate() || @@ -89,7 +93,7 @@ int main(int argc, const char** argv) { LOG_DEBUG("%s", ctx_params.to_string().c_str()); LOG_DEBUG("%s", default_gen_params.to_string().c_str()); - sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(false); + sd_ctx_params_t sd_ctx_params = ctx_params.to_sd_ctx_params_t(); ServerModels models; if (!ctx_params.model_path.empty() || !ctx_params.diffusion_model_path.empty()) { const auto started = std::chrono::steady_clock::now(); diff --git a/examples/server/model_catalog.cpp b/examples/server/model_catalog.cpp index a3c59206d2..f665567096 100644 --- a/examples/server/model_catalog.cpp +++ b/examples/server/model_catalog.cpp @@ -19,7 +19,7 @@ struct ModelSlot { const char* id; const char* label; std::vector categories; - std::string SDContextParams::*field; + std::string SDContextParams::* field; }; static const ModelSlot k_model_slots[] = { @@ -32,7 +32,7 @@ static const ModelSlot k_model_slots[] = { {"clip_vision", "CLIP Vision", {"clip_vision"}, &SDContextParams::clip_vision_path}, {"t5xxl", "T5", {"text_encoders"}, &SDContextParams::t5xxl_path}, {"llm", "LLM Text Encoder", {"text_encoders"}, &SDContextParams::llm_path}, - {"llm_vision", "LLM Vision Encoder", {"clip_vision"}, &SDContextParams::llm_vision_path}, + {"llm_vision", "LLM Vision Encoder", {"text_encoders", "clip_vision"}, &SDContextParams::llm_vision_path}, {"tokenizer", "Tokenizer", {"tokenizers"}, &SDContextParams::tokenizer}, {"vae", "VAE", {"vae"}, &SDContextParams::vae_path}, {"audio_vae", "Audio VAE", {"vae"}, &SDContextParams::audio_vae_path}, @@ -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(); @@ -273,7 +281,7 @@ static void apply_recipe_defaults(const json& defaults, SDGenerationParams& para } static void apply_recipe_context(const json& settings, SDContextParams& params) { - check_keys(settings, {"vae_format", "model_args"}); + check_keys(settings, {"vae_format", "model_args", "backend", "params_backend"}); if (settings.contains("vae_format")) { const auto format = settings.at("vae_format").get(); if (format != "auto" && format != "flux" && format != "sd3" && format != "flux2" && format != "wan") @@ -283,8 +291,15 @@ static void apply_recipe_context(const json& settings, SDContextParams& params) if (settings.contains("model_args")) { params.model_args = settings.at("model_args").get(); } + if (settings.contains("backend")) { + params.backend = settings.at("backend").get(); + } + if (settings.contains("params_backend")) { + params.params_backend = settings.at("params_backend").get(); + } } +// Known recipe feature names with their defaults; recipes may flip any of these. static const std::map k_recipe_features = { {"img2img", true}, {"mask_image", true}, @@ -306,6 +321,9 @@ static const std::map k_recipe_features = { {"init_image", false}, {"end_image", false}, {"control_frames", false}, + {"ref_videos", false}, + {"ref_audios", false}, + {"ref_video_audios", false}, {"vace", false}, {"high_noise", false}, {"cache_unet", false}, @@ -359,17 +377,44 @@ json model_generation_features(const ServerRuntime& runtime) { if (runtime.models->active_recipe.value("source", "") == "command_line") return features; // Optional components only become usable once selected and loaded. - features["ip_adapter"] = features.value("ip_adapter", false) && !runtime.ctx_params->ip_adapter_path.empty(); - features["controlnet"] = features.value("controlnet", false) && !runtime.ctx_params->control_net_path.empty(); - features["high_noise"] = features.value("high_noise", false) && !runtime.ctx_params->high_noise_diffusion_model_path.empty(); - features["preview_tae"] = features.value("preview_tae", false) && !runtime.ctx_params->taesd_path.empty(); + features["ip_adapter"] = features.value("ip_adapter", false) && !runtime.ctx_params->ip_adapter_path.empty(); + features["controlnet"] = features.value("controlnet", false) && !runtime.ctx_params->control_net_path.empty(); + features["high_noise"] = features.value("high_noise", false) && !runtime.ctx_params->high_noise_diffusion_model_path.empty(); + features["preview_tae"] = features.value("preview_tae", false) && !runtime.ctx_params->taesd_path.empty(); + features["ref_audios"] = features.value("ref_audios", false) && + (!runtime.ctx_params->audio_vae_path.empty() || !runtime.ctx_params->audio_encoder_path.empty()); + features["ref_video_audios"] = features.value("ref_video_audios", false) && !runtime.ctx_params->audio_vae_path.empty(); if (!runtime.ctx_params->pulid_weights_path.empty()) features["slg"] = false; return features; } +static void validate_video_constraints(const json& constraints) { + check_keys(constraints, {"frame_step", "frame_offset", "frame_rounding", "min_frames", "max_frames", "fixed_fps"}); + for (const auto* key : {"frame_step", "frame_offset", "min_frames", "max_frames", "fixed_fps"}) { + const int minimum = std::string(key) == "frame_offset" ? 0 : 1; + const int maximum = std::string(key) == "fixed_fps" ? 240 : 4096; + if (constraints.contains(key) && + (!constraints[key].is_number_integer() || constraints[key] < minimum || constraints[key] > maximum)) { + throw std::runtime_error(std::string("video_constraints.") + key + " is outside its integer range"); + } + } + const int step = constraints.value("frame_step", 1); + const int offset = constraints.value("frame_offset", 0); + const int minimum = constraints.value("min_frames", 1); + const int maximum = constraints.value("max_frames", 4096); + if (offset >= step || minimum > maximum || + offset + std::max(0, (minimum - offset + step - 1) / step) * step > maximum) { + throw std::runtime_error("video_constraints must describe a nonempty frame range with frame_offset < frame_step"); + } + if (constraints.contains("frame_rounding") && + constraints["frame_rounding"] != "floor" && constraints["frame_rounding"] != "ceil") { + throw std::runtime_error("video_constraints.frame_rounding must be floor or ceil"); + } +} + 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"}); + 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()) || recipe.at("name").get().empty()) { throw std::runtime_error("invalid schema_version, id or name"); @@ -377,6 +422,13 @@ static void validate_recipe(const json& recipe) { if (recipe.contains("description") && !recipe["description"].is_string()) { throw std::runtime_error("description must be a string"); } + if (recipe.contains("size_multiple") && + (!recipe["size_multiple"].is_number_integer() || recipe["size_multiple"] < 1 || recipe["size_multiple"] > 4096)) { + throw std::runtime_error("size_multiple must be an integer from 1 to 4096"); + } + if (recipe.contains("video_constraints")) { + validate_video_constraints(recipe["video_constraints"]); + } if (recipe.contains("model_url")) { const auto url = recipe.at("model_url").get(); const std::string prefix = "https://huggingface.co/"; @@ -405,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 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().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) @@ -429,6 +494,7 @@ static void validate_recipe(const json& recipe) { static json normalize_recipe(json recipe) { validate_recipe(recipe); + recipe["size_multiple"] = recipe.value("size_multiple", 8); recipe["effective_features"] = resolve_recipe_features(recipe); for (auto& slot : recipe.at("slots")) { const auto& spec = *find_model_slot(slot.at("id").get()); @@ -592,6 +658,7 @@ void initialize_server_models(ServerRuntime& runtime) { {"description", "Uses the loaded command-line model and its generation defaults."}, {"slots", json::array()}, {"defaults", json::object()}, + {"size_multiple", 8}, {"features", features}, }; models.active_recipe["effective_features"] = resolve_recipe_features(models.active_recipe); @@ -600,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(); - 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; @@ -618,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 allowed; for (const auto& slot : recipe.at("slots")) { const auto id = slot.at("id").get(); @@ -629,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()); @@ -652,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().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()).empty(); diff --git a/examples/server/model_catalog.h b/examples/server/model_catalog.h index 7d49a19241..2e45fd183f 100644 --- a/examples/server/model_catalog.h +++ b/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 diff --git a/examples/server/openai_edit_policy.h b/examples/server/openai_edit_policy.h deleted file mode 100644 index c963a887f8..0000000000 --- a/examples/server/openai_edit_policy.h +++ /dev/null @@ -1,44 +0,0 @@ -#ifndef __SD_EXAMPLES_SERVER_OPENAI_EDIT_POLICY_H__ -#define __SD_EXAMPLES_SERVER_OPENAI_EDIT_POLICY_H__ - -#include "common/resource_owners.hpp" - -enum class OpenAIEditInitImageAction { - KeepExisting, - UseUploaded, - Clear, -}; - -inline bool openai_edit_uses_uploaded_image_as_init(bool supports_ref_image_conditioning, - bool has_ref_images, - bool has_mask, - bool has_explicit_strength_override) { - return !supports_ref_image_conditioning || !has_ref_images || has_mask || has_explicit_strength_override; -} - -inline OpenAIEditInitImageAction apply_openai_edit_init_image_policy( - SDImageOwner& init_image, - SDImageOwner& uploaded_image, - bool uploaded_image_was_first_reference, - bool supports_ref_image_conditioning, - bool has_ref_images, - bool has_mask, - bool has_extra_init_image_override, - bool has_explicit_strength_override) { - if (!uploaded_image_was_first_reference || has_extra_init_image_override) { - return OpenAIEditInitImageAction::KeepExisting; - } - - if (openai_edit_uses_uploaded_image_as_init(supports_ref_image_conditioning, - has_ref_images, - has_mask, - has_explicit_strength_override)) { - init_image = std::move(uploaded_image); - return OpenAIEditInitImageAction::UseUploaded; - } - - init_image.reset(); - return OpenAIEditInitImageAction::Clear; -} - -#endif // __SD_EXAMPLES_SERVER_OPENAI_EDIT_POLICY_H__ diff --git a/examples/server/recipes/README.md b/examples/server/recipes/README.md index 5043f6d1b8..dcf9d69ab3 100644 --- a/examples/server/recipes/README.md +++ b/examples/server/recipes/README.md @@ -77,3 +77,35 @@ These are built-in catalog rules. Users may keep personal presets in `/recipes/`; they need not add those presets to the built-in list. See [the recipe schema and feature reference](../api.md#user-recipe-format) for supported fields. + +## Output dimension alignment + +Set `size_multiple` when the runtime's output dimension alignment differs from +the default of `8`. Match `GenerationRequest::align_image_size` in +[`src/pipeline/request.cpp`](../../../src/pipeline/request.cpp): the VAE scale +factor multiplied by the model definition's `diffusion_down_factor`. The scale +comes from [`VAE::get_scale_factor`](../../../src/model/vae/vae.hpp), with the +PiD override in [`diffusion_engine.cpp`](../../../src/pipeline/diffusion_engine.cpp); +model geometry is defined under [`src/model/integrations`](../../../src/model/integrations/). + +Use the runtime's output alignment, not the nominal patch size: some models pad +patches internally and crop their output back to the requested size. Pixel-space +models with an alignment of `1` explicitly use `size_multiple: 1`. MiniT2I also +has this alignment, but its separate fixed `512 × 512` output restriction still +applies. Reference-image preprocessing constraints are independent of this field. + +## Video length constraints + +Video recipes declare `video_constraints` to match `align_video_frames` in +[`diffusion_engine.cpp`](../../../src/pipeline/diffusion_engine.cpp) and the FPS +and AnimateDiff overrides in +[`request.cpp`](../../../src/pipeline/request.cpp). Wan, HunyuanVideo and LingBot +round down to `4n + 1`; LTX rounds down to `8n + 1`. MiniMax-H3 rounds up to +`17n + 5`, with at least 5 frames and fixed 24 FPS. AnimateDiff accepts whole +frames up to 32. Wan 2.2 S2V fixes FPS at 16. + +Declare fixed FPS only where the runtime enforces it; recommended FPS belongs +in `defaults.fps`. These constraints guide UI conversion and estimates, not +runtime behavior. Audio-driven output and temporal upscaling may still change +the final frame count. See the [recipe schema](../api.md#user-recipe-format) +for field defaults and validation rules. diff --git a/examples/server/recipes/animatediff-sd15.json b/examples/server/recipes/animatediff-sd15.json index a7c8338f19..e7bf07d3a8 100644 --- a/examples/server/recipes/animatediff-sd15.json +++ b/examples/server/recipes/animatediff-sd15.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "animatediff-sd15", "name": "AnimateDiff SD 1.5", + "size_multiple": 64, "description": "Use Video Generation with an SD 1.5 checkpoint and compatible motion module.", "model_url": "https://huggingface.co/guoyww/animatediff", "slots": [ @@ -22,6 +23,13 @@ "required": false } ], + "video_constraints": { + "frame_step": 1, + "frame_offset": 0, + "frame_rounding": "floor", + "min_frames": 1, + "max_frames": 32 + }, "defaults": { "sampling_method": "euler", "video_frames": 8, diff --git a/examples/server/recipes/chroma-radiance.json b/examples/server/recipes/chroma-radiance.json index 6a0be8a562..2d35f434cc 100644 --- a/examples/server/recipes/chroma-radiance.json +++ b/examples/server/recipes/chroma-radiance.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "chroma-radiance", "name": "Chroma Radiance", + "size_multiple": 1, "description": "Pixel-space generation without a VAE. Also accepts the x0 and experimental patch-32 variants.", "model_url": "https://huggingface.co/lodestones/Chroma1-Radiance", "slots": [ diff --git a/examples/server/recipes/cosxl-edit.json b/examples/server/recipes/cosxl-edit.json index 49407fccd1..eb794248ae 100644 --- a/examples/server/recipes/cosxl-edit.json +++ b/examples/server/recipes/cosxl-edit.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "cosxl-edit", "name": "CosXL Edit", + "size_multiple": 64, "description": "Supply the image to edit as a reference image.", "model_url": "https://huggingface.co/stabilityai/cosxl", "slots": [ diff --git a/examples/server/recipes/defaults.md b/examples/server/recipes/defaults.md index 2be4dfbf68..efe596c92b 100644 --- a/examples/server/recipes/defaults.md +++ b/examples/server/recipes/defaults.md @@ -5,6 +5,9 @@ Reviewed against model-publisher cards and reference implementations on 2026-10- These are starting values, not limits on supported resolutions or generation settings. Built-in step defaults above 20 are reduced to 20; the cited sources may recommend more steps. Users can still choose a higher step count. +Except for PiD recipes, built-in dimensions above 1,048,576 pixels (1024 x 1024) are reduced to 1024 x 1024, even when the cited source uses a larger size. +PiD recipes retain their 2048 x 2048 defaults. +Users can still choose a higher resolution. When several aspect ratios or ranges are published, the selected example or square bucket is noted below. A dash means no recipe override: the available primary sources did not establish a fixed value, or the size depends on the input image. Ordinary DiT recipes explicitly use Euler; PiD retains its LCM sampling path. @@ -27,7 +30,7 @@ For community edit adapters and other optional extensions, these are the base-mo | `flux1-canny` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) | | `flux1-depth` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) | | `flux1-dev` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-dev) | -| `flux1-fill` | 1 | 20 | 1232 × 1632 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev) | +| `flux1-fill` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev) | | `flux1-kontext` | 1 | 20 | — | [Source](https://raw.githubusercontent.com/huggingface/diffusers/main/src/diffusers/pipelines/flux/pipeline_flux_kontext.py) Output size normally follows the reference image; no fixed size is imposed by this recipe. | | `flux1-pulid` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-dev) FLUX.1-dev defaults; identity conditioning is provided by PuLID. | | `flux1-schnell` | 1 | 4 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.1-schnell) | @@ -38,7 +41,7 @@ For community edit adapters and other optional extensions, these are the base-mo | `flux2-klein-base-9b` | 4 | 20 | 1024 × 1024 | [Source](https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B) | | `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 | 2048 × 2048 | [Source](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev) | +| `hidream-o1-dev` | 1 | 20 | 1024 × 1024 | [Source](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev) | | `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. | @@ -46,10 +49,10 @@ For community edit adapters and other optional extensions, these are the base-mo | `iris-3b` | 3 | 20 | 1024 × 1024 | [Source](https://huggingface.co/speridlabs/iris-3b) | | `krea2-raw` | 3.5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/krea/Krea-2-Raw) Base model defaults; edit requires a matching community LoRA with its own settings. | | `krea2-raw-edit` | 3.5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/krea/Krea-2-Raw) Base model defaults; edit requires a matching community LoRA with its own settings. | -| `krea2-turbo` | 1 | 8 | 2048 × 2048 | [Source](https://huggingface.co/krea/Krea-2-Turbo) Official CLI example; edit requires a matching community LoRA. | -| `krea2-turbo-edit` | 1 | 8 | 2048 × 2048 | [Source](https://huggingface.co/krea/Krea-2-Turbo) Official CLI example; edit requires a matching community LoRA. | -| `lens` | 5 | 20 | 1440 × 1440 | [Source](https://huggingface.co/microsoft/Lens) Official model card: RL model, default base_resolution 1440. | -| `lens-turbo` | 1 | 4 | 1440 × 1440 | [Source](https://huggingface.co/microsoft/Lens-Turbo) Official model card: 4-step distilled model, default base_resolution 1440. | +| `krea2-turbo` | 1 | 8 | 1024 × 1024 | [Source](https://huggingface.co/krea/Krea-2-Turbo) Official CLI example; edit requires a matching community LoRA. | +| `krea2-turbo-edit` | 1 | 8 | 1024 × 1024 | [Source](https://huggingface.co/krea/Krea-2-Turbo) Official CLI example; edit requires a matching community LoRA. | +| `lens` | 5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/microsoft/Lens) Official model card: RL model, default base_resolution 1440. | +| `lens-turbo` | 1 | 4 | 1024 × 1024 | [Source](https://huggingface.co/microsoft/Lens-Turbo) Official model card: 4-step distilled model, default base_resolution 1440. | | `lingbot-video-dense` | 3 | 20 | 832 × 480 | [Source](https://huggingface.co/robbyant/lingbot-video-dense-1.3b) Official first-stage inference example; optional refiner is separate. | | `lingbot-video-moe` | 3 | 20 | 832 × 480 | [Source](https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b) Official inference example. | | `llada-image` | 5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/inclusionAI/LLaDA-Image) | @@ -62,12 +65,13 @@ For community edit adapters and other optional extensions, these are the base-mo | `mage-flow` | 5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/microsoft/Mage-Flow) | | `mage-flow-edit` | 5 | 20 | 1024 × 1024 | [Source](https://huggingface.co/microsoft/Mage-Flow) | | `mage-flow-turbo` | 1 | 4 | 1024 × 1024 | [Source](https://huggingface.co/microsoft/Mage-Flow) | -| `ming-image` | 1 | 12 | 2048 × 2048 | [Source](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design) | +| `ming-image` | 1 | 12 | 1024 × 1024 | [Source](https://huggingface.co/inclusionAI/Ming-Image-0.1-Design) | | `minimax-h3-fl2va` | 1 | — | 1344 × 768 | [Source](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3) Guidance-distilled; official 768p canvas. No fixed step count is prescribed here; the upstream API counts sigma grid points rather than sd.cpp denoising steps. | | `minimax-h3-ref2va` | 1 | — | 1344 × 768 | [Source](https://huggingface.co/docs/diffusers/main/en/api/pipelines/minimax_h3) Guidance-distilled; official 768p canvas. No fixed step count is prescribed here; the upstream API counts sigma grid points rather than sd.cpp denoising steps. | | `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. | @@ -81,9 +85,9 @@ For community edit adapters and other optional extensions, these are the base-mo | `pixart-alpha-1024` | 4.5 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/huggingface/diffusers/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py) | | `pixart-alpha-512` | 4.5 | 20 | 512 × 512 | [Source](https://raw.githubusercontent.com/huggingface/diffusers/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_alpha.py) | | `pixart-sigma` | 4.5 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/huggingface/diffusers/main/src/diffusers/pipelines/pixart_alpha/pipeline_pixart_sigma.py) | -| `qwen-image` | 4 | 20 | 1328 × 1328 | [Source](https://huggingface.co/Qwen/Qwen-Image) Official square resolution bucket; true_cfg_scale maps to cfg_scale. | -| `qwen-image-21` | 1 | 20 | 2048 × 2048 | [Source](https://raw.githubusercontent.com/QwenLM/Qwen-Image-2.1/main/README.md) Original checkpoint defaults; select the matching TAEQI2.1 to replace its VAE. | -| `qwen-image-21-turbo` | 1 | 8 | 2048 × 2048 | [Model card](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo), [fixed sigmas](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/blob/main/model_index.json). Standalone Turbo weights; the fixed 8-step schedule is applied on model load. | +| `qwen-image` | 4 | 20 | 1024 × 1024 | [Source](https://huggingface.co/Qwen/Qwen-Image) The source uses a larger square resolution bucket; true_cfg_scale maps to cfg_scale. | +| `qwen-image-21` | 1 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/QwenLM/Qwen-Image-2.1/main/README.md) Original checkpoint sampling defaults; select the matching TAEQI2.1 to replace its VAE. | +| `qwen-image-21-turbo` | 1 | 8 | 1024 × 1024 | [Model card](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo), [fixed sigmas](https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo/blob/main/model_index.json). Standalone Turbo weights; the fixed 8-step schedule is applied on model load. | | `qwen-image-edit` | 4 | 20 | — | [Source](https://huggingface.co/Qwen/Qwen-Image-Edit) Official editing example; output dimensions follow reference-image aspect ratio. | | `qwen-image-edit-2509` | 4 | 20 | — | [Source](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) Official editing example; output dimensions follow reference-image aspect ratio. | | `qwen-image-edit-2511` | 4 | 20 | — | [Source](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) Official editing example; output dimensions follow reference-image aspect ratio. | @@ -107,7 +111,7 @@ For community edit adapters and other optional extensions, these are the base-mo | `sdxl-turbo` | 1 | 1 | 512 × 512 | [Source](https://huggingface.co/stabilityai/sdxl-turbo) | | `sdxs` | 1 | 1 | 512 × 512 | [Source](https://huggingface.co/IDKiro/sdxs-512-dreamshaper) | | `sefi-image` | 4 | 20 | — | [Source](https://huggingface.co/SeFi-Image/SeFi-Image-5B-Base) Base/RL family defaults; Turbo checkpoints require 4 steps / CFG 1. | -| `sensenova-u15` | 4 | 20 | 2048 × 2048 | [Source](https://raw.githubusercontent.com/OpenSenseNova/SenseNova-U1/main/examples/README.md) | +| `sensenova-u15` | 4 | 20 | 1024 × 1024 | [Source](https://raw.githubusercontent.com/OpenSenseNova/SenseNova-U1/main/examples/README.md) | | `wan21-flf2v` | 5 | 20 | 1280 × 720 | [Source](https://raw.githubusercontent.com/Wan-Video/Wan2.1/main/generate.py) | | `wan21-i2v` | 5 | 20 | 832 × 480 | [Source](https://raw.githubusercontent.com/Wan-Video/Wan2.1/main/generate.py) 480p base defaults; also accepts compatible SkyReels V2 and single Wan 2.2 I2V experts. | | `wan21-t2v` | 5 | 20 | 1280 × 720 | [Source](https://raw.githubusercontent.com/Wan-Video/Wan2.1/main/generate.py) 14B base defaults; single Wan 2.2 experts and acceleration LoRAs need matching settings. | diff --git a/examples/server/recipes/ernie-image-turbo.json b/examples/server/recipes/ernie-image-turbo.json index 5702f5efcb..90ec37df55 100644 --- a/examples/server/recipes/ernie-image-turbo.json +++ b/examples/server/recipes/ernie-image-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ernie-image-turbo", "name": "ERNIE Image Turbo", + "size_multiple": 16, "model_url": "https://huggingface.co/baidu/ERNIE-Image-Turbo", "slots": [ { diff --git a/examples/server/recipes/ernie-image.json b/examples/server/recipes/ernie-image.json index a610000083..6373fed518 100644 --- a/examples/server/recipes/ernie-image.json +++ b/examples/server/recipes/ernie-image.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ernie-image", "name": "ERNIE Image", + "size_multiple": 16, "model_url": "https://huggingface.co/baidu/ERNIE-Image", "slots": [ { diff --git a/examples/server/recipes/fibo2-dev.json b/examples/server/recipes/fibo2-dev.json new file mode 100644 index 0000000000..061b9395e0 --- /dev/null +++ b/examples/server/recipes/fibo2-dev.json @@ -0,0 +1,40 @@ +{ + "schema_version": 1, + "id": "fibo2-dev", + "name": "Fibo 2.0 Dev", + "size_multiple": 16, + "description": "Text-to-image generation and native reference-image editing. Select a separate vision projector for editing when it is not bundled in the text encoder.", + "model_url": "https://huggingface.co/briaai/Fibo-2.0-dev", + "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 for separate GGUF text encoders)" + }, + { + "id": "vae", + "required": true, + "label": "FLUX.2 VAE" + } + ], + "defaults": { + "sampling_method": "euler", + "cfg_scale": 5, + "steps": 20, + "width": 1024, + "height": 1024 + }, + "features": [ + "flow_shift", + "ref_images" + ] +} diff --git a/examples/server/recipes/fibo2-edit-turbo.json b/examples/server/recipes/fibo2-edit-turbo.json new file mode 100644 index 0000000000..a78f5d3678 --- /dev/null +++ b/examples/server/recipes/fibo2-edit-turbo.json @@ -0,0 +1,40 @@ +{ + "schema_version": 1, + "id": "fibo2-edit-turbo", + "name": "Fibo 2.0 Edit Turbo", + "size_multiple": 16, + "description": "Dedicated reference-image editing checkpoint; supply reference images. Select a separate vision projector when it is not bundled in the text encoder.", + "model_url": "https://huggingface.co/briaai/Fibo-2.0-edit-turbo", + "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 for separate GGUF text encoders)" + }, + { + "id": "vae", + "required": true, + "label": "FLUX.2 VAE" + } + ], + "defaults": { + "sampling_method": "euler", + "cfg_scale": 1, + "steps": 4, + "width": 1024, + "height": 1024 + }, + "features": [ + "flow_shift", + "ref_images" + ] +} diff --git a/examples/server/recipes/flux1-fill.json b/examples/server/recipes/flux1-fill.json index 5e6485ad10..fe27a48e48 100644 --- a/examples/server/recipes/flux1-fill.json +++ b/examples/server/recipes/flux1-fill.json @@ -39,8 +39,8 @@ "scheduler": "flux", "cfg_scale": 1, "steps": 20, - "width": 1232, - "height": 1632 + "width": 1024, + "height": 1024 }, "features": [ "clip_skip", diff --git a/examples/server/recipes/flux2-dev.json b/examples/server/recipes/flux2-dev.json index 4d27ea16b6..12cf6af75c 100644 --- a/examples/server/recipes/flux2-dev.json +++ b/examples/server/recipes/flux2-dev.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux2-dev", "name": "FLUX.2 dev", + "size_multiple": 16, "description": "Supports text-to-image and reference-image editing.", "model_url": "https://huggingface.co/black-forest-labs/FLUX.2-dev", "slots": [ diff --git a/examples/server/recipes/flux2-klein-4b.json b/examples/server/recipes/flux2-klein-4b.json index f09360cfbc..67281e1e7c 100644 --- a/examples/server/recipes/flux2-klein-4b.json +++ b/examples/server/recipes/flux2-klein-4b.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux2-klein-4b", "name": "FLUX.2 Klein 4B", + "size_multiple": 16, "description": "Supports reference-image editing. The FLUX.2 small decoder is an alternative VAE.", "model_url": "https://huggingface.co/black-forest-labs/FLUX.2-klein-4B", "slots": [ diff --git a/examples/server/recipes/flux2-klein-9b.json b/examples/server/recipes/flux2-klein-9b.json index ad10d6e516..1c8d50f659 100644 --- a/examples/server/recipes/flux2-klein-9b.json +++ b/examples/server/recipes/flux2-klein-9b.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux2-klein-9b", "name": "FLUX.2 Klein 9B", + "size_multiple": 16, "description": "Supports text-to-image and reference-image editing.", "model_url": "https://huggingface.co/black-forest-labs/FLUX.2-klein-9B", "slots": [ diff --git a/examples/server/recipes/flux2-klein-base-4b.json b/examples/server/recipes/flux2-klein-base-4b.json index 56c69f9801..3fff7d9680 100644 --- a/examples/server/recipes/flux2-klein-base-4b.json +++ b/examples/server/recipes/flux2-klein-base-4b.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux2-klein-base-4b", "name": "FLUX.2 Klein Base 4B", + "size_multiple": 16, "description": "Supports text-to-image and reference-image editing.", "model_url": "https://huggingface.co/black-forest-labs/FLUX.2-klein-base-4B", "slots": [ diff --git a/examples/server/recipes/flux2-klein-base-9b.json b/examples/server/recipes/flux2-klein-base-9b.json index e9bd53c6cf..b9a9eb6dd7 100644 --- a/examples/server/recipes/flux2-klein-base-9b.json +++ b/examples/server/recipes/flux2-klein-base-9b.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux2-klein-base-9b", "name": "FLUX.2 Klein Base 9B", + "size_multiple": 16, "description": "Use the 9B base transformer with Qwen 3 8B. Supports reference-image editing.", "model_url": "https://huggingface.co/black-forest-labs/FLUX.2-klein-base-9B", "slots": [ diff --git a/examples/server/recipes/flux3-klein-base.json b/examples/server/recipes/flux3-klein-base.json index 29f6f68196..06789c2e58 100644 --- a/examples/server/recipes/flux3-klein-base.json +++ b/examples/server/recipes/flux3-klein-base.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux3-klein-base", "name": "FLUX.3 Image Klein Base", + "size_multiple": 16, "model_url": "https://huggingface.co/black-forest-labs/flux-3-image-klein-base", "slots": [ { diff --git a/examples/server/recipes/flux3-klein.json b/examples/server/recipes/flux3-klein.json index 164bce12ad..e63314dd1c 100644 --- a/examples/server/recipes/flux3-klein.json +++ b/examples/server/recipes/flux3-klein.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "flux3-klein", "name": "FLUX.3 Image Klein", + "size_multiple": 16, "description": "Supports text-to-image and reference-image editing.", "model_url": "https://huggingface.co/black-forest-labs/flux-3-image-klein", "slots": [ diff --git a/examples/server/recipes/hidream-o1-dev.json b/examples/server/recipes/hidream-o1-dev.json index 8e95f5ceac..b5c35921ae 100644 --- a/examples/server/recipes/hidream-o1-dev.json +++ b/examples/server/recipes/hidream-o1-dev.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "hidream-o1-dev", "name": "HiDream O1 Image dev", + "size_multiple": 1, "description": "Bundled checkpoint; also supports reference-image editing.", "model_url": "https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev", "slots": [ @@ -13,8 +14,8 @@ "defaults": { "cfg_scale": 1, "steps": 20, - "width": 2048, - "height": 2048, + "width": 1024, + "height": 1024, "sampling_method": "euler" }, "features": [ diff --git a/examples/server/recipes/hunyuan-video-15-i2v.json b/examples/server/recipes/hunyuan-video-15-i2v.json index a4d672f7b1..5dc0509f91 100644 --- a/examples/server/recipes/hunyuan-video-15-i2v.json +++ b/examples/server/recipes/hunyuan-video-15-i2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "hunyuan-video-15-i2v", "name": "HunyuanVideo 1.5 I2V", + "size_multiple": 16, "description": "Select HunyuanVideo 1.5 I2V weights and supply a start image in Video Generation.", "model_url": "https://huggingface.co/tencent/HunyuanVideo-1.5", "slots": [ @@ -36,6 +37,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/hunyuan-video-15-t2v.json b/examples/server/recipes/hunyuan-video-15-t2v.json index d477b8063d..81aaceee64 100644 --- a/examples/server/recipes/hunyuan-video-15-t2v.json +++ b/examples/server/recipes/hunyuan-video-15-t2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "hunyuan-video-15-t2v", "name": "HunyuanVideo 1.5 T2V", + "size_multiple": 16, "description": "Select HunyuanVideo 1.5 T2V weights and use Video Generation. Use the separate I2V recipe for I2V weights.", "model_url": "https://huggingface.co/tencent/HunyuanVideo-1.5", "slots": [ @@ -36,6 +37,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/ideogram4.json b/examples/server/recipes/ideogram4.json index 2d5230019b..133cf5f359 100644 --- a/examples/server/recipes/ideogram4.json +++ b/examples/server/recipes/ideogram4.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ideogram4", "name": "Ideogram 4", + "size_multiple": 16, "description": "Select the matching unconditional transformer for classifier-free guidance.", "model_url": "https://huggingface.co/ideogram-ai/ideogram-4-fp8", "slots": [ diff --git a/examples/server/recipes/instruct-pix2pix.json b/examples/server/recipes/instruct-pix2pix.json index e819509108..e1f69a8a48 100644 --- a/examples/server/recipes/instruct-pix2pix.json +++ b/examples/server/recipes/instruct-pix2pix.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "instruct-pix2pix", "name": "Instruct-Pix2Pix", + "size_multiple": 64, "description": "Supply the image to edit as a reference image.", "model_url": "https://huggingface.co/timbrooks/instruct-pix2pix", "slots": [ diff --git a/examples/server/recipes/iris-3b.json b/examples/server/recipes/iris-3b.json index c1b5d84fcd..96a617efe5 100644 --- a/examples/server/recipes/iris-3b.json +++ b/examples/server/recipes/iris-3b.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "iris-3b", "name": "Iris 3B", + "size_multiple": 16, "description": "Pixel-space model; no VAE is required.", "model_url": "https://huggingface.co/speridlabs/iris-3b", "slots": [ diff --git a/examples/server/recipes/krea2-turbo-edit.json b/examples/server/recipes/krea2-turbo-edit.json index fa248c083b..e8c521d354 100644 --- a/examples/server/recipes/krea2-turbo-edit.json +++ b/examples/server/recipes/krea2-turbo-edit.json @@ -39,8 +39,8 @@ "sampling_method": "euler", "cfg_scale": 1, "steps": 8, - "width": 2048, - "height": 2048 + "width": 1024, + "height": 1024 }, "features": [ "flow_shift", diff --git a/examples/server/recipes/krea2-turbo.json b/examples/server/recipes/krea2-turbo.json index 23dad12fd5..f19513f77a 100644 --- a/examples/server/recipes/krea2-turbo.json +++ b/examples/server/recipes/krea2-turbo.json @@ -32,8 +32,8 @@ "defaults": { "cfg_scale": 1, "steps": 8, - "width": 2048, - "height": 2048, + "width": 1024, + "height": 1024, "sampling_method": "euler" }, "features": [ diff --git a/examples/server/recipes/lens-turbo.json b/examples/server/recipes/lens-turbo.json index e9282eafc4..364320e7b8 100644 --- a/examples/server/recipes/lens-turbo.json +++ b/examples/server/recipes/lens-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "lens-turbo", "name": "LENS Turbo", + "size_multiple": 16, "model_url": "https://huggingface.co/microsoft/Lens-Turbo", "slots": [ { @@ -36,8 +37,8 @@ "defaults": { "cfg_scale": 1, "steps": 4, - "width": 1440, - "height": 1440, + "width": 1024, + "height": 1024, "sampling_method": "euler" }, "features": [ diff --git a/examples/server/recipes/lens.json b/examples/server/recipes/lens.json index 0a89542293..0e049c82f9 100644 --- a/examples/server/recipes/lens.json +++ b/examples/server/recipes/lens.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "lens", "name": "LENS", + "size_multiple": 16, "model_url": "https://huggingface.co/microsoft/Lens", "slots": [ { @@ -36,8 +37,8 @@ "defaults": { "cfg_scale": 5, "steps": 20, - "width": 1440, - "height": 1440, + "width": 1024, + "height": 1024, "sampling_method": "euler" }, "features": [ diff --git a/examples/server/recipes/lingbot-video-dense.json b/examples/server/recipes/lingbot-video-dense.json index 5f313a0cd9..3affb48516 100644 --- a/examples/server/recipes/lingbot-video-dense.json +++ b/examples/server/recipes/lingbot-video-dense.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "lingbot-video-dense", "name": "LingBot Video Dense 1.3B", + "size_multiple": 16, "description": "Use Video Generation.", "model_url": "https://huggingface.co/robbyant/lingbot-video-dense-1.3b", "slots": [ @@ -30,6 +31,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "video_frames": 33, "cfg_scale": 3, diff --git a/examples/server/recipes/lingbot-video-moe.json b/examples/server/recipes/lingbot-video-moe.json index cdbce61ee9..55bf59f0e3 100644 --- a/examples/server/recipes/lingbot-video-moe.json +++ b/examples/server/recipes/lingbot-video-moe.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "lingbot-video-moe", "name": "LingBot Video MoE 30B A3B", + "size_multiple": 16, "description": "Use Video Generation.", "model_url": "https://huggingface.co/robbyant/lingbot-video-moe-30b-a3b", "slots": [ @@ -30,6 +31,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "video_frames": 33, "cfg_scale": 3, diff --git a/examples/server/recipes/llada-image-turbo.json b/examples/server/recipes/llada-image-turbo.json index 0589bb3846..4965ddc990 100644 --- a/examples/server/recipes/llada-image-turbo.json +++ b/examples/server/recipes/llada-image-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "llada-image-turbo", "name": "LLaDA Image Turbo", + "size_multiple": 16, "description": "Use matching turbo text encoder and connectors. Editing needs connectors that include the SigVQ encoder.", "model_url": "https://huggingface.co/inclusionAI/LLaDA-Image-Turbo", "slots": [ diff --git a/examples/server/recipes/llada-image.json b/examples/server/recipes/llada-image.json index 27887c849e..e16feda062 100644 --- a/examples/server/recipes/llada-image.json +++ b/examples/server/recipes/llada-image.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "llada-image", "name": "LLaDA Image (base)", + "size_multiple": 16, "description": "Use matching base transformer, text encoder and connectors; do not mix them with Turbo components. Editing needs connectors with SigVQ and 1024x1024 output.", "model_url": "https://huggingface.co/inclusionAI/LLaDA-Image", "slots": [ diff --git a/examples/server/recipes/ltx23-dev.json b/examples/server/recipes/ltx23-dev.json index 83bb934e86..f0bf3f2883 100644 --- a/examples/server/recipes/ltx23-dev.json +++ b/examples/server/recipes/ltx23-dev.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ltx23-dev", "name": "LTX 2.3 dev", + "size_multiple": 32, "description": "Supports text, start-image and start/end-image video. Audio VAE is optional for silent video. Configure latent upscaling separately in generation settings. Optional distillation LoRAs use this recipe; select the LoRA and adjust the sampling settings separately.", "model_url": "https://huggingface.co/Lightricks/LTX-2.3", "slots": [ @@ -39,6 +40,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 8, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/ltx23-distilled.json b/examples/server/recipes/ltx23-distilled.json index a2a6ac33c2..3e4d08e9db 100644 --- a/examples/server/recipes/ltx23-distilled.json +++ b/examples/server/recipes/ltx23-distilled.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ltx23-distilled", "name": "LTX 2.3 Distilled 1.1", + "size_multiple": 32, "description": "Supports text, start-image and start/end-image video. Audio VAE is optional for silent video. Configure latent upscaling separately in generation settings.", "model_url": "https://huggingface.co/Lightricks/LTX-2.3", "slots": [ @@ -39,6 +40,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 8, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 121, diff --git a/examples/server/recipes/ltx25-dev.json b/examples/server/recipes/ltx25-dev.json index 47efbc76fd..011a49edaa 100644 --- a/examples/server/recipes/ltx25-dev.json +++ b/examples/server/recipes/ltx25-dev.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ltx25-dev", "name": "LTX 2.5 dev", + "size_multiple": 32, "description": "Use the conv video VAE and LTX-specific Gemma 4 with bundled projection. Audio VAE is optional for silent video.", "model_url": "https://huggingface.co/Lightricks/LTX-2.5", "slots": [ @@ -35,6 +36,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 8, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "fps": 24, diff --git a/examples/server/recipes/mage-flow-edit.json b/examples/server/recipes/mage-flow-edit.json index 82d66f6bd2..e0c552457c 100644 --- a/examples/server/recipes/mage-flow-edit.json +++ b/examples/server/recipes/mage-flow-edit.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "mage-flow-edit", "name": "Mage Flow Edit", + "size_multiple": 16, "description": "Supply reference images. Select a separate vision projector when it is not bundled in the text encoder.", "model_url": "https://huggingface.co/microsoft/Mage-Flow-Edit", "slots": [ diff --git a/examples/server/recipes/mage-flow-turbo.json b/examples/server/recipes/mage-flow-turbo.json index 5cf9bb25a7..8e8a9d63d4 100644 --- a/examples/server/recipes/mage-flow-turbo.json +++ b/examples/server/recipes/mage-flow-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "mage-flow-turbo", "name": "Mage Flow Turbo", + "size_multiple": 16, "model_url": "https://huggingface.co/microsoft/Mage-Flow-Turbo", "slots": [ { diff --git a/examples/server/recipes/mage-flow.json b/examples/server/recipes/mage-flow.json index 4cbd15934c..c2751f086a 100644 --- a/examples/server/recipes/mage-flow.json +++ b/examples/server/recipes/mage-flow.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "mage-flow", "name": "Mage Flow", + "size_multiple": 16, "model_url": "https://huggingface.co/microsoft/Mage-Flow", "slots": [ { diff --git a/examples/server/recipes/ming-image.json b/examples/server/recipes/ming-image.json index 6c368e8e1a..f100581b4b 100644 --- a/examples/server/recipes/ming-image.json +++ b/examples/server/recipes/ming-image.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "ming-image", "name": "Ming Image 0.1", + "size_multiple": 16, "model_url": "https://huggingface.co/inclusionAI/Ming-Image-0.1-Design", "slots": [ { @@ -36,8 +37,8 @@ "defaults": { "cfg_scale": 1, "steps": 12, - "width": 2048, - "height": 2048, + "width": 1024, + "height": 1024, "sampling_method": "euler" }, "features": [ diff --git a/examples/server/recipes/minimax-h3-fl2va.json b/examples/server/recipes/minimax-h3-fl2va.json index 3971923917..f0af22bd9d 100644 --- a/examples/server/recipes/minimax-h3-fl2va.json +++ b/examples/server/recipes/minimax-h3-fl2va.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "minimax-h3-fl2va", "name": "MiniMax H3 FL2VA", + "size_multiple": 32, "description": "Generate video and audio, optionally with start/end images. Supports full and pruned FL2VA weights. Audio VAE is optional for silent video.", "model_url": "https://huggingface.co/MiniMaxAI/MiniMax-H3", "slots": [ @@ -35,6 +36,13 @@ "tae" ] ], + "video_constraints": { + "frame_step": 17, + "frame_offset": 5, + "frame_rounding": "ceil", + "min_frames": 5, + "fixed_fps": 24 + }, "defaults": { "video_frames": 56, "fps": 24, diff --git a/examples/server/recipes/minimax-h3-ref2va.json b/examples/server/recipes/minimax-h3-ref2va.json index 085f5d5fed..9ee5f63d7d 100644 --- a/examples/server/recipes/minimax-h3-ref2va.json +++ b/examples/server/recipes/minimax-h3-ref2va.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "minimax-h3-ref2va", "name": "MiniMax H3 Ref2VA", + "size_multiple": 32, "description": "Supply image, video or audio references with compatible Ref2VA weights. Audio VAE is optional for silent video.", "model_url": "https://huggingface.co/MiniMaxAI/MiniMax-H3", "slots": [ @@ -35,6 +36,13 @@ "tae" ] ], + "video_constraints": { + "frame_step": 17, + "frame_offset": 5, + "frame_rounding": "ceil", + "min_frames": 5, + "fixed_fps": 24 + }, "defaults": { "video_frames": 56, "fps": 24, @@ -45,7 +53,10 @@ }, "features": [ "flow_shift", - "ref_images" + "ref_images", + "ref_videos", + "ref_audios", + "ref_video_audios" ], "disabled_features": [ "img2img", diff --git a/examples/server/recipes/minit2i.json b/examples/server/recipes/minit2i.json index 2c7b94eba6..a82333d939 100644 --- a/examples/server/recipes/minit2i.json +++ b/examples/server/recipes/minit2i.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "minit2i", "name": "MiniT2I", + "size_multiple": 1, "description": "Pixel-space model; no VAE is required.", "model_url": "https://huggingface.co/MiniT2I/MiniT2I", "slots": [ diff --git a/examples/server/recipes/moss-image-base.json b/examples/server/recipes/moss-image-base.json index bc2cc7d9b1..273eeea7b4 100644 --- a/examples/server/recipes/moss-image-base.json +++ b/examples/server/recipes/moss-image-base.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "moss-image-base", "name": "MOSS Image Base", + "size_multiple": 16, "description": "Select the MOSS-Image-Base transformer with its matching text encoder and FLUX.2 VAE or TAE.", "model_url": "https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base", "slots": [ diff --git a/examples/server/recipes/moss-image-sft.json b/examples/server/recipes/moss-image-sft.json index e0239b0bcd..d63dd2eb27 100644 --- a/examples/server/recipes/moss-image-sft.json +++ b/examples/server/recipes/moss-image-sft.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "moss-image-sft", "name": "MOSS Image SFT", + "size_multiple": 16, "description": "Select the MOSS-Image-SFT transformer with its matching text encoder and FLUX.2 VAE or TAE.", "model_url": "https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT", "slots": [ diff --git a/examples/server/recipes/nitrosd-realism.json b/examples/server/recipes/nitrosd-realism.json index 75f2d43ff8..ae7345db72 100644 --- a/examples/server/recipes/nitrosd-realism.json +++ b/examples/server/recipes/nitrosd-realism.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "nitrosd-realism", "name": "NitroSD Realism", + "size_multiple": 64, "model_url": "https://huggingface.co/ChenDY/NitroFusion", "slots": [ { diff --git a/examples/server/recipes/p-h3-image.json b/examples/server/recipes/p-h3-image.json new file mode 100644 index 0000000000..1098bab9eb --- /dev/null +++ b/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" + ] +} diff --git a/examples/server/recipes/pid-15-flux2.json b/examples/server/recipes/pid-15-flux2.json index 4fb3db3a72..84cc941c51 100644 --- a/examples/server/recipes/pid-15-flux2.json +++ b/examples/server/recipes/pid-15-flux2.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-15-flux2", "name": "PiD 1.5 (Flux.2 VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a matching Flux.2 PiD checkpoint and VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid-15-qwen-image.json b/examples/server/recipes/pid-15-qwen-image.json index 9db9e7422b..1415ac755d 100644 --- a/examples/server/recipes/pid-15-qwen-image.json +++ b/examples/server/recipes/pid-15-qwen-image.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-15-qwen-image", "name": "PiD 1.5 (Qwen Image VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a matching Qwen Image PiD checkpoint and VAE. Use the Qwen Image 2D VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid-15-sd3.json b/examples/server/recipes/pid-15-sd3.json index 3980b86710..c92e9c0388 100644 --- a/examples/server/recipes/pid-15-sd3.json +++ b/examples/server/recipes/pid-15-sd3.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-15-sd3", "name": "PiD 1.5 (SD3 VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a matching SD3 PiD checkpoint and VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid-15.json b/examples/server/recipes/pid-15.json index f7913bee8e..e98d762cb7 100644 --- a/examples/server/recipes/pid-15.json +++ b/examples/server/recipes/pid-15.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-15", "name": "PiD 1.5 (Flux / Z-Image VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a Flux/Z-Image PiD checkpoint and matching Flux VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid-flux2.json b/examples/server/recipes/pid-flux2.json index 9c802dc153..827b4f2ab1 100644 --- a/examples/server/recipes/pid-flux2.json +++ b/examples/server/recipes/pid-flux2.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-flux2", "name": "PiD (Flux.2 VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a matching Flux.2 PiD checkpoint and VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid-qwen-image.json b/examples/server/recipes/pid-qwen-image.json index 4503b1cd69..2155260280 100644 --- a/examples/server/recipes/pid-qwen-image.json +++ b/examples/server/recipes/pid-qwen-image.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-qwen-image", "name": "PiD (Qwen Image VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a matching Qwen Image PiD checkpoint and VAE. Use the Qwen Image 2D VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid-sd3.json b/examples/server/recipes/pid-sd3.json index 6dad2ca09e..2170b8fb1c 100644 --- a/examples/server/recipes/pid-sd3.json +++ b/examples/server/recipes/pid-sd3.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid-sd3", "name": "PiD (SD3 VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a matching SD3 PiD checkpoint and VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pid.json b/examples/server/recipes/pid.json index 2479449659..77c02e969a 100644 --- a/examples/server/recipes/pid.json +++ b/examples/server/recipes/pid.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pid", "name": "PiD (Flux / Z-Image VAE)", + "size_multiple": 1, "description": "Supply a reference image; the first image is used as the source latent. Select a Flux/Z-Image PiD checkpoint and matching Flux VAE.", "model_url": "https://huggingface.co/nvidia/PiD", "slots": [ diff --git a/examples/server/recipes/pixart-alpha-1024.json b/examples/server/recipes/pixart-alpha-1024.json index 1144f87b2a..e17bf54a0c 100644 --- a/examples/server/recipes/pixart-alpha-1024.json +++ b/examples/server/recipes/pixart-alpha-1024.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pixart-alpha-1024", "name": "PixArt Alpha XL 1024", + "size_multiple": 16, "description": "Use the matching PixArt Alpha transformer and SD 1.x-style VAE. Resolution/aspect-ratio micro-conditioning is not currently applied by the runtime.", "model_url": "https://huggingface.co/PixArt-alpha/PixArt-XL-2-1024-MS", "slots": [ diff --git a/examples/server/recipes/pixart-alpha-512.json b/examples/server/recipes/pixart-alpha-512.json index 804a7a0c9e..7f6817d022 100644 --- a/examples/server/recipes/pixart-alpha-512.json +++ b/examples/server/recipes/pixart-alpha-512.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pixart-alpha-512", "name": "PixArt Alpha XL 512", + "size_multiple": 16, "description": "Use the matching PixArt Alpha transformer and SD 1.x-style VAE. Resolution/aspect-ratio micro-conditioning is not currently applied by the runtime.", "model_url": "https://huggingface.co/PixArt-alpha/PixArt-XL-2-512x512", "slots": [ diff --git a/examples/server/recipes/pixart-sigma.json b/examples/server/recipes/pixart-sigma.json index ed4d2f0dec..dd79065ce4 100644 --- a/examples/server/recipes/pixart-sigma.json +++ b/examples/server/recipes/pixart-sigma.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "pixart-sigma", "name": "PixArt Sigma XL 1024", + "size_multiple": 16, "model_url": "https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", "slots": [ { diff --git a/examples/server/recipes/qwen-image-21-turbo.json b/examples/server/recipes/qwen-image-21-turbo.json index 2ef9fa183e..5962b92a3b 100644 --- a/examples/server/recipes/qwen-image-21-turbo.json +++ b/examples/server/recipes/qwen-image-21-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "qwen-image-21-turbo", "name": "Qwen Image 2.1 Turbo", + "size_multiple": 32, "description": "Use the standalone Turbo diffusion weights for generation and editing with the fixed 8-step schedule. Select the Qwen Image 2.1 VAE or TAEQI2.1, and supply a vision projector when needed.", "model_url": "https://huggingface.co/Qwen/Qwen-Image-2.1-Turbo", "slots": [ @@ -40,8 +41,8 @@ "sampling_method": "euler", "cfg_scale": 1, "steps": 8, - "width": 2048, - "height": 2048, + "width": 1024, + "height": 1024, "custom_sigmas": [ 1.0, 0.978453, diff --git a/examples/server/recipes/qwen-image-21.json b/examples/server/recipes/qwen-image-21.json index 4b6c67b86e..12c907d868 100644 --- a/examples/server/recipes/qwen-image-21.json +++ b/examples/server/recipes/qwen-image-21.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "qwen-image-21", "name": "Qwen Image 2.1", + "size_multiple": 32, "description": "Supports generation and editing. Select the full Qwen Image 2.1 VAE or TAEQI2.1. Supply a separate vision projector when needed. Select acceleration LoRAs and adjust their settings separately.", "model_url": "https://huggingface.co/Qwen/Qwen-Image-2.1", "slots": [ @@ -40,8 +41,8 @@ "sampling_method": "euler", "cfg_scale": 1, "steps": 20, - "width": 2048, - "height": 2048 + "width": 1024, + "height": 1024 }, "features": [ "flow_shift", diff --git a/examples/server/recipes/qwen-image.json b/examples/server/recipes/qwen-image.json index 417a621ae2..468f8f2613 100644 --- a/examples/server/recipes/qwen-image.json +++ b/examples/server/recipes/qwen-image.json @@ -35,8 +35,8 @@ "flow_shift": 3.0, "cfg_scale": 4, "steps": 20, - "width": 1328, - "height": 1328 + "width": 1024, + "height": 1024 }, "features": [ "flow_shift" diff --git a/examples/server/recipes/sd-turbo.json b/examples/server/recipes/sd-turbo.json index 7723c81071..cd5955266e 100644 --- a/examples/server/recipes/sd-turbo.json +++ b/examples/server/recipes/sd-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd-turbo", "name": "SD Turbo", + "size_multiple": 64, "description": "Select a single-file SD-Turbo checkpoint. Uses one step with CFG disabled (scale 1 in sd.cpp); use more steps for partial image-to-image denoising.", "model_url": "https://huggingface.co/stabilityai/sd-turbo", "slots": [ diff --git a/examples/server/recipes/sd1-checkpoint.json b/examples/server/recipes/sd1-checkpoint.json index 38d2386775..240bdf0957 100644 --- a/examples/server/recipes/sd1-checkpoint.json +++ b/examples/server/recipes/sd1-checkpoint.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd1-checkpoint", "name": "SD 1.x", + "size_multiple": 64, "description": "Supports SD 1.x, compatible LCM checkpoints and layer-count variants. For LCM weights, select LCM sampling and adjust steps/CFG. Tiny SD Diffusers directories must first be converted into a single-file checkpoint.", "model_url": "https://huggingface.co/CompVis/stable-diffusion-v1-4", "slots": [ diff --git a/examples/server/recipes/sd1-controlnet.json b/examples/server/recipes/sd1-controlnet.json index eeaee22fb3..ddd764940c 100644 --- a/examples/server/recipes/sd1-controlnet.json +++ b/examples/server/recipes/sd1-controlnet.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd1-controlnet", "name": "SD 1.5 ControlNet", + "size_multiple": 64, "description": "Supply a control image; input image and mask are optional for image-to-image or inpainting.", "model_url": "https://huggingface.co/lllyasviel/control_v11p_sd15_canny", "slots": [ diff --git a/examples/server/recipes/sd1-inpainting.json b/examples/server/recipes/sd1-inpainting.json index f9f3064f3f..a4c0aa2c87 100644 --- a/examples/server/recipes/sd1-inpainting.json +++ b/examples/server/recipes/sd1-inpainting.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd1-inpainting", "name": "SD 1.5 Inpainting", + "size_multiple": 64, "description": "Select an inpainting checkpoint, then supply an input image and mask.", "slots": [ { diff --git a/examples/server/recipes/sd1-ip-adapter.json b/examples/server/recipes/sd1-ip-adapter.json index 6c0ab1a6b0..7e14c15668 100644 --- a/examples/server/recipes/sd1-ip-adapter.json +++ b/examples/server/recipes/sd1-ip-adapter.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd1-ip-adapter", "name": "SD 1.5 IP-Adapter / Plus", + "size_multiple": 64, "description": "Select matching SD 1.5 IP-Adapter or Plus weights and CLIP Vision H, then supply an IP-Adapter image.", "model_url": "https://huggingface.co/h94/IP-Adapter", "slots": [ diff --git a/examples/server/recipes/sd2-512-checkpoint.json b/examples/server/recipes/sd2-512-checkpoint.json index 77838c2805..e5df5af305 100644 --- a/examples/server/recipes/sd2-512-checkpoint.json +++ b/examples/server/recipes/sd2-512-checkpoint.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd2-512-checkpoint", "name": "SD 2.1 512", + "size_multiple": 64, "description": "Supports SD 2.x 512 and compatible layer-count variants. Tiny SD2 Diffusers directories must first be converted into a single-file checkpoint.", "slots": [ { diff --git a/examples/server/recipes/sd2-768-checkpoint.json b/examples/server/recipes/sd2-768-checkpoint.json index 808b187d7a..8e04418569 100644 --- a/examples/server/recipes/sd2-768-checkpoint.json +++ b/examples/server/recipes/sd2-768-checkpoint.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd2-768-checkpoint", "name": "SD 2.1 768", + "size_multiple": 64, "slots": [ { "id": "model", diff --git a/examples/server/recipes/sd2-inpainting.json b/examples/server/recipes/sd2-inpainting.json index 4980f7f0a5..7e5e31262d 100644 --- a/examples/server/recipes/sd2-inpainting.json +++ b/examples/server/recipes/sd2-inpainting.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sd2-inpainting", "name": "SD 2 Inpainting", + "size_multiple": 64, "description": "Select an inpainting checkpoint, then supply Init Image and Mask Image.", "slots": [ { diff --git a/examples/server/recipes/sdxl-checkpoint.json b/examples/server/recipes/sdxl-checkpoint.json index 372d5a6d4a..a1fea1bd27 100644 --- a/examples/server/recipes/sdxl-checkpoint.json +++ b/examples/server/recipes/sdxl-checkpoint.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxl-checkpoint", "name": "SDXL", + "size_multiple": 64, "description": "Supports SDXL and compatible reduced-depth checkpoints such as SSD-1B and Segmind Vega.", "model_url": "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0", "slots": [ diff --git a/examples/server/recipes/sdxl-controlnet.json b/examples/server/recipes/sdxl-controlnet.json index c50efbd2b4..3ffb9b4c83 100644 --- a/examples/server/recipes/sdxl-controlnet.json +++ b/examples/server/recipes/sdxl-controlnet.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxl-controlnet", "name": "SDXL ControlNet", + "size_multiple": 64, "description": "Select a compatible SDXL ControlNet and supply a preprocessed control image. Input image and mask are optional for image-to-image or inpainting.", "model_url": "https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0", "slots": [ diff --git a/examples/server/recipes/sdxl-inpainting.json b/examples/server/recipes/sdxl-inpainting.json index 574423d08f..f4ef972175 100644 --- a/examples/server/recipes/sdxl-inpainting.json +++ b/examples/server/recipes/sdxl-inpainting.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxl-inpainting", "name": "SDXL Inpainting", + "size_multiple": 64, "description": "Select an inpainting checkpoint, then supply Init Image and Mask Image.", "model_url": "https://huggingface.co/diffusers/stable-diffusion-xl-1.0-inpainting-0.1", "slots": [ diff --git a/examples/server/recipes/sdxl-ip-adapter.json b/examples/server/recipes/sdxl-ip-adapter.json index b649fb1d5a..f0bc8cdafb 100644 --- a/examples/server/recipes/sdxl-ip-adapter.json +++ b/examples/server/recipes/sdxl-ip-adapter.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxl-ip-adapter", "name": "SDXL IP-Adapter / Plus", + "size_multiple": 64, "description": "Select matching SDXL ViT-H IP-Adapter or Plus weights and CLIP Vision H, then supply an IP-Adapter image.", "model_url": "https://huggingface.co/h94/IP-Adapter", "slots": [ diff --git a/examples/server/recipes/sdxl-photomaker.json b/examples/server/recipes/sdxl-photomaker.json index 025fb7668e..c8b5790256 100644 --- a/examples/server/recipes/sdxl-photomaker.json +++ b/examples/server/recipes/sdxl-photomaker.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxl-photomaker", "name": "SDXL PhotoMaker", + "size_multiple": 64, "description": "Select PhotoMaker v1 and supply identity images. Include the img trigger in the prompt.", "model_url": "https://huggingface.co/TencentARC/PhotoMaker", "slots": [ diff --git a/examples/server/recipes/sdxl-turbo.json b/examples/server/recipes/sdxl-turbo.json index 8650e6c581..496291d71a 100644 --- a/examples/server/recipes/sdxl-turbo.json +++ b/examples/server/recipes/sdxl-turbo.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxl-turbo", "name": "SDXL Turbo", + "size_multiple": 64, "model_url": "https://huggingface.co/stabilityai/sdxl-turbo", "slots": [ { diff --git a/examples/server/recipes/sdxs.json b/examples/server/recipes/sdxs.json index 19e218ac5c..f83b060a8a 100644 --- a/examples/server/recipes/sdxs.json +++ b/examples/server/recipes/sdxs.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sdxs", "name": "SDXS", + "size_multiple": 64, "model_url": "https://huggingface.co/IDKiro/sdxs-512-dreamshaper", "slots": [ { diff --git a/examples/server/recipes/sefi-image.json b/examples/server/recipes/sefi-image.json index 35e933360d..168809d1ea 100644 --- a/examples/server/recipes/sefi-image.json +++ b/examples/server/recipes/sefi-image.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sefi-image", "name": "SeFi Image 5B RL", + "size_multiple": 16, "model_url": "https://huggingface.co/SeFi-Image/SeFi-Image-5B-Base", "slots": [ { diff --git a/examples/server/recipes/sensenova-u15.json b/examples/server/recipes/sensenova-u15.json index afd20466b6..9ca386e446 100644 --- a/examples/server/recipes/sensenova-u15.json +++ b/examples/server/recipes/sensenova-u15.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "sensenova-u15", "name": "SenseNova U1.5", + "size_multiple": 32, "description": "Bundled MoT checkpoint.", "model_url": "https://huggingface.co/sensenova/SenseNova-U1.5-8B-MoT", "slots": [ @@ -14,8 +15,8 @@ "flow_shift": 3.0, "cfg_scale": 4, "steps": 20, - "width": 2048, - "height": 2048, + "width": 1024, + "height": 1024, "sampling_method": "euler" }, "features": [ diff --git a/examples/server/recipes/wan21-flf2v.json b/examples/server/recipes/wan21-flf2v.json index 02f0e80a91..6bb4c33d62 100644 --- a/examples/server/recipes/wan21-flf2v.json +++ b/examples/server/recipes/wan21-flf2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan21-flf2v", "name": "Wan 2.1 FLF2V 14B", + "size_multiple": 16, "description": "Select FLF2V weights and supply both start and end images.", "model_url": "https://huggingface.co/Wan-AI/Wan2.1-FLF2V-14B-720P", "slots": [ @@ -35,6 +36,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "flow_shift": 16, diff --git a/examples/server/recipes/wan21-i2v.json b/examples/server/recipes/wan21-i2v.json index d8d8d13ee1..32d5fc7ace 100644 --- a/examples/server/recipes/wan21-i2v.json +++ b/examples/server/recipes/wan21-i2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan21-i2v", "name": "Wan 2.1 I2V", + "size_multiple": 16, "description": "Image-to-video with Wan 2.1, compatible SkyReels V2 weights, or a single Wan 2.2 I2V expert. CLIP Vision is required by Wan 2.1 I2V and is unused by Wan 2.2. Defaults target Wan 2.1 480p.", "model_url": "https://huggingface.co/Wan-AI/Wan2.1-I2V-14B-480P", "slots": [ @@ -35,6 +36,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/wan21-t2v.json b/examples/server/recipes/wan21-t2v.json index 367611c413..998078cb07 100644 --- a/examples/server/recipes/wan21-t2v.json +++ b/examples/server/recipes/wan21-t2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan21-t2v", "name": "Wan 2.1 T2V", + "size_multiple": 16, "description": "Text-to-video with Wan 2.1 or a compatible single Wan 2.2 T2V expert. Defaults target Wan 2.1 14B; adjust settings for the selected weights and acceleration LoRAs.", "model_url": "https://huggingface.co/Wan-AI/Wan2.1-T2V-14B", "slots": [ @@ -30,6 +31,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/wan21-vace.json b/examples/server/recipes/wan21-vace.json index 7b7f7e41b8..d7f834380b 100644 --- a/examples/server/recipes/wan21-vace.json +++ b/examples/server/recipes/wan21-vace.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan21-vace", "name": "Wan 2.1 VACE 1.3B / 14B", + "size_multiple": 16, "description": "Select matching VACE weights. Supports text, reference images and control video through Video Generation.", "model_url": "https://huggingface.co/Wan-AI/Wan2.1-VACE-14B", "slots": [ @@ -30,6 +31,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/wan22-i2v.json b/examples/server/recipes/wan22-i2v.json index 62afdfce4f..c7ab2ed8db 100644 --- a/examples/server/recipes/wan22-i2v.json +++ b/examples/server/recipes/wan22-i2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan22-i2v", "name": "Wan 2.2 I2V A14B", + "size_multiple": 16, "description": "Image-to-video. Select a start image in Video Generation after loading. LightX2V acceleration LoRAs must match the I2V low- and high-noise models; configure LoRAs and sampling separately.", "model_url": "https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B", "slots": [ @@ -35,6 +36,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/wan22-s2v.json b/examples/server/recipes/wan22-s2v.json index ba1ee4f7c3..641ed6e59c 100644 --- a/examples/server/recipes/wan22-s2v.json +++ b/examples/server/recipes/wan22-s2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan22-s2v", "name": "Wan 2.2 S2V 14B", + "size_multiple": 16, "description": "Supply a start image and reference audio through Video Generation.", "model_url": "https://huggingface.co/Wan-AI/Wan2.2-S2V-14B", "slots": [ @@ -35,6 +36,13 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1, + "fixed_fps": 16 + }, "defaults": { "sampling_method": "euler", "flow_shift": 3, @@ -44,7 +52,8 @@ }, "features": [ "flow_shift", - "init_image" + "init_image", + "ref_audios" ], "disabled_features": [ "img2img", diff --git a/examples/server/recipes/wan22-t2v.json b/examples/server/recipes/wan22-t2v.json index 905bcfa4ff..fba962fcb2 100644 --- a/examples/server/recipes/wan22-t2v.json +++ b/examples/server/recipes/wan22-t2v.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan22-t2v", "name": "Wan 2.2 T2V A14B", + "size_multiple": 16, "description": "Text-to-video. Use Video Generation after loading. LightX2V acceleration LoRAs must match the T2V low- and high-noise models; configure LoRAs and sampling separately.", "model_url": "https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B", "slots": [ @@ -35,6 +36,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/wan22-ti2v-5b.json b/examples/server/recipes/wan22-ti2v-5b.json index 3e2aa35b88..bc433611a6 100644 --- a/examples/server/recipes/wan22-ti2v-5b.json +++ b/examples/server/recipes/wan22-ti2v-5b.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "wan22-ti2v-5b", "name": "Wan 2.2 TI2V 5B", + "size_multiple": 32, "description": "Use Wan 2.2 VAE. Supports text, start-image and start/end-image conditioning.", "model_url": "https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B", "slots": [ @@ -30,6 +31,12 @@ "tae" ] ], + "video_constraints": { + "frame_step": 4, + "frame_offset": 1, + "frame_rounding": "floor", + "min_frames": 1 + }, "defaults": { "sampling_method": "euler", "video_frames": 33, diff --git a/examples/server/recipes/z-image-l2p.json b/examples/server/recipes/z-image-l2p.json index 0375f1387b..192d547266 100644 --- a/examples/server/recipes/z-image-l2p.json +++ b/examples/server/recipes/z-image-l2p.json @@ -2,6 +2,7 @@ "schema_version": 1, "id": "z-image-l2p", "name": "Z-Image L2P", + "size_multiple": 16, "description": "Pixel-space L2P checkpoint with Qwen 3 4B; no VAE is needed. Use dimensions divisible by 16.", "model_url": "https://huggingface.co/zhen-nan/L2P", "slots": [ diff --git a/examples/server/routes_models.cpp b/examples/server/routes_models.cpp index ea14eb821a..f33057758c 100644 --- a/examples/server/routes_models.cpp +++ b/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; @@ -94,7 +112,7 @@ static void change_model(ServerRuntime& runtime, const httplib::Request& req, ht *runtime.ctx_params = unload ? models.base_context : std::move(context); *runtime.default_gen_params = unload ? models.base_defaults : std::move(defaults); if (!unload) { - auto params = runtime.ctx_params->to_sd_ctx_params_t(false); + auto params = runtime.ctx_params->to_sd_ctx_params_t(); models.context.reset(new_sd_ctx(¶ms)); if (!models.context) { throw std::runtime_error("Model initialization failed; check the server log for component compatibility and memory errors"); diff --git a/examples/server/routes_openai.cpp b/examples/server/routes_openai.cpp index 29dfe313a7..f721e1495d 100644 --- a/examples/server/routes_openai.cpp +++ b/examples/server/routes_openai.cpp @@ -7,7 +7,6 @@ #include "common/common.h" #include "common/media_io.h" #include "common/resource_owners.hpp" -#include "openai_edit_policy.h" static std::string extract_and_remove_sd_cpp_extra_args(std::string& text) { std::regex re("(.*?)"); @@ -75,7 +74,7 @@ static bool build_openai_generation_request(const httplib::Request& req, error_message = "invalid params"; return false; } - return true; + return validate_server_img_gen_request(runtime, request, error_message); } static bool build_openai_edit_request(const httplib::Request& req, @@ -158,49 +157,24 @@ static bool build_openai_edit_request(const httplib::Request& req, request.gen_params.height = height; request.gen_params.batch_count = n; - std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt); - const json extra_args = sd_cpp_extra_args_str.empty() - ? json::object() - : json::parse(sd_cpp_extra_args_str, nullptr, false); - const bool has_extra_init_image_override = extra_args.is_object() && extra_args.contains("init_image"); - const bool has_extra_ref_images_override = extra_args.is_object() && extra_args.contains("ref_images"); - const bool has_extra_strength_override = extra_args.is_object() && - extra_args.contains("strength") && - extra_args["strength"].is_number(); - const bool has_extra_mask_image = extra_args.is_object() && - extra_args.contains("mask_image") && - !extra_args["mask_image"].is_null(); - const bool supports_ref_image_conditioning = sd_ctx_supports_ref_image_conditioning(runtime.sd_ctx); - const bool may_need_uploaded_init_image = !has_extra_init_image_override && - (!supports_ref_image_conditioning || !mask_bytes.empty() || - has_extra_mask_image || has_extra_strength_override || has_extra_ref_images_override); - SDImageOwner uploaded_image_for_init; - bool uploaded_image_was_first_reference = false; + std::string sd_cpp_extra_args_str = extract_and_remove_sd_cpp_extra_args(request.gen_params.prompt); for (auto& bytes : images_bytes) { int img_w = 0; int img_h = 0; int resolved_channel = 0; uint8_t* raw_pixels = load_image_from_memory( - reinterpret_cast(bytes.data()), - static_cast(bytes.size()), - img_w, img_h, resolved_channel, - 0, 0, - 0); + reinterpret_cast(bytes.data()), + static_cast(bytes.size()), + img_w, img_h, resolved_channel, + 0, 0, + 0); if (raw_pixels == nullptr) { continue; } - const bool is_first_ref_image = request.gen_params.ref_images.empty(); SDImageOwner image_owner({(uint32_t)img_w, (uint32_t)img_h, (uint32_t)resolved_channel, raw_pixels}); request.gen_params.set_width_and_height_if_unset(image_owner.get().width, image_owner.get().height); - if (is_first_ref_image) { - uploaded_image_was_first_reference = true; - if (may_need_uploaded_init_image) { - uploaded_image_for_init = image_owner; - } - } - request.gen_params.ref_images.push_back(std::move(image_owner)); } @@ -231,28 +205,13 @@ static bool build_openai_edit_request(const httplib::Request& req, return false; } - const auto init_image_action = apply_openai_edit_init_image_policy( - request.gen_params.init_image, - uploaded_image_for_init, - uploaded_image_was_first_reference, - supports_ref_image_conditioning, - !request.gen_params.ref_images.empty(), - request.gen_params.mask_image.get().data != nullptr, - has_extra_init_image_override, - has_extra_strength_override); - if (init_image_action == OpenAIEditInitImageAction::UseUploaded && - request.gen_params.init_image.get().data == nullptr) { - error_message = "could not allocate init image"; - return false; - } - // Intentionally disable prompt-embedded LoRA tag parsing for server APIs. if (!request.gen_params.resolve_and_validate(IMG_GEN, "", runtime.ctx_params->hires_upscalers_dir, true)) { error_message = "invalid params"; return false; } - return true; + return validate_server_img_gen_request(runtime, request, error_message); } static bool execute_sync_img_gen_request(ServerRuntime& runtime, @@ -264,12 +223,14 @@ static bool execute_sync_img_gen_request(ServerRuntime& runtime, { std::lock_guard lock(*runtime.sd_ctx_mutex); - sd_image_t* raw_results = nullptr; - if (!generate_image(runtime.sd_ctx, &img_gen_params, &raw_results, &num_results)) { - raw_results = nullptr; - num_results = 0; + sd_image_generation_result_t* generated = nullptr; + generate_image(runtime.sd_ctx, &img_gen_params, &generated); + results.adopt(generated, runtime.ctx_params->lora_model_dir); + num_results = results.count(); + if (!results.empty() && !apply_server_adetailer(runtime, request, img_gen_params, results)) { + error_message = "ADetailer failed"; + return false; } - results.adopt(raw_results, num_results); } if (results.empty()) { @@ -322,35 +283,34 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { out["data"] = json::array(); out["output_format"] = request.output_format; - int result_count = results.count(); - int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1; + int result_count = results.count(); + for (int i = 0; i < result_count; ++i) { if (results[i].data == nullptr) { continue; } std::string params = request.gen_params.embed_image_metadata - ? get_image_params(*runtime->ctx_params, - request.gen_params, - request.gen_params.seed + i / images_per_batch) + ? results.generation_info(i) : ""; auto image_bytes = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : request.output_format == "webp" ? EncodedImageFormat::WEBP : EncodedImageFormat::PNG, - results[i].data, - results[i].width, - results[i].height, - results[i].channel, - params, - request.output_compression); + results[i].data, + results[i].width, + results[i].height, + results[i].channel, + params, + request.output_compression); if (image_bytes.empty()) { LOG_ERROR("write image to mem failed"); continue; } json item; - item["b64_json"] = base64_encode(image_bytes); + item["b64_json"] = base64_encode(image_bytes); + item["generation_info"] = json::parse(results.generation_info(i)); out["data"].push_back(item); } @@ -396,26 +356,25 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { out["data"] = json::array(); out["output_format"] = request.output_format; - int result_count = results.count(); - int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, result_count / request.gen_params.batch_count) : 1; + int result_count = results.count(); + for (int i = 0; i < result_count; ++i) { if (results[i].data == nullptr) { continue; } std::string params = request.gen_params.embed_image_metadata - ? get_image_params(*runtime->ctx_params, - request.gen_params, - request.gen_params.seed + i / images_per_batch) + ? results.generation_info(i) : ""; auto image_bytes = encode_image_to_vector(request.output_format == "jpeg" ? EncodedImageFormat::JPEG : EncodedImageFormat::PNG, - results[i].data, - results[i].width, - results[i].height, - results[i].channel, - params, - request.output_compression); + results[i].data, + results[i].width, + results[i].height, + results[i].channel, + params, + request.output_compression); json item; - item["b64_json"] = base64_encode(image_bytes); + item["b64_json"] = base64_encode(image_bytes); + item["generation_info"] = json::parse(results.generation_info(i)); out["data"].push_back(item); } diff --git a/examples/server/routes_sdapi.cpp b/examples/server/routes_sdapi.cpp index 4b8e7a5b1e..0d37898a8d 100644 --- a/examples/server/routes_sdapi.cpp +++ b/examples/server/routes_sdapi.cpp @@ -3,9 +3,12 @@ #include #include #include +#include #include +#include #include #include +#include #include "common/common.h" #include "common/media_io.h" @@ -81,6 +84,202 @@ static enum sample_method_t get_sdapi_sample_method(std::string name) { return it != hardcoded.end() ? it->second : SAMPLE_METHOD_COUNT; } +static bool parse_sdapi_adetailer(const json& body, + ServerRuntime& runtime, + bool img2img, + ImgGenJobRequest& request, + std::string& error_message) { + request.adetailer_enabled = false; + request.adetailer_passes.clear(); + if (!body.contains("alwayson_scripts")) { + return true; + } + try { + const auto& scripts = body.at("alwayson_scripts"); + if (!scripts.is_object()) { + throw std::invalid_argument("alwayson_scripts must be an object"); + } + if (!scripts.contains("ADetailer")) { + return true; + } + const auto& args = scripts.at("ADetailer").at("args"); + if (!args.is_array()) { + throw std::invalid_argument("ADetailer.args must be an array"); + } + if (args.empty() || (args[0].is_boolean() && !args[0].get())) { + return true; + } + static const std::unordered_set keys = { + "ad_model", "ad_model_classes", "ad_tab_enable", "ad_prompt", "ad_negative_prompt", + "ad_confidence", "ad_mask_filter_method", "ad_mask_k", "ad_mask_min_ratio", "ad_mask_max_ratio", + "ad_dilate_erode", "ad_x_offset", "ad_y_offset", "ad_mask_merge_invert", "ad_mask_blur", + "ad_denoising_strength", "ad_inpaint_only_masked", "ad_inpaint_only_masked_padding", + "ad_use_inpaint_width_height", "ad_inpaint_width", "ad_inpaint_height", "ad_use_steps", "ad_steps", + "ad_use_cfg_scale", "ad_cfg_scale", "ad_use_sampler", "ad_sampler", "ad_scheduler", + "ad_use_checkpoint", "ad_checkpoint", "ad_use_vae", "ad_vae", "ad_use_noise_multiplier", + "ad_noise_multiplier", "ad_use_clip_skip", "ad_clip_skip", "ad_restore_face", "ad_controlnet_model", + "ad_controlnet_module", "ad_controlnet_weight", "ad_controlnet_guidance_start", "ad_controlnet_guidance_end", "is_api"}; + for (const auto& config : args) { + if (!config.is_object()) { + continue; + } + auto boolean = [&](const char* key, bool fallback) { + if (!config.contains(key)) + return fallback; + if (!config.at(key).is_boolean()) + throw std::invalid_argument(std::string(key) + " must be boolean"); + return config.at(key).get(); + }; + auto text = [&](const char* key, const char* fallback) { + if (!config.contains(key)) + return std::string(fallback); + if (!config.at(key).is_string()) + throw std::invalid_argument(std::string(key) + " must be a string"); + return config.at(key).get(); + }; + auto number = [&](const char* key, double fallback, double minimum, double maximum) { + double value = fallback; + if (config.contains(key)) { + if (!config.at(key).is_number()) + throw std::invalid_argument(std::string(key) + " must be numeric"); + value = config.at(key).get(); + } + if (!std::isfinite(value) || value < minimum || value > maximum) + throw std::invalid_argument(std::string(key) + " is out of range"); + return value; + }; + auto integer = [&](const char* key, int fallback, int minimum, int maximum = std::numeric_limits::max()) { + if (config.contains(key) && !config.at(key).is_number_integer()) + throw std::invalid_argument(std::string(key) + " must be an integer"); + return static_cast(number(key, fallback, minimum, maximum)); + }; + const std::string model = text("ad_model", "None"); + if (model == "None" || !boolean("ad_tab_enable", true)) { + continue; + } + for (const auto& item : config.items()) { + if (keys.count(item.key()) == 0) + throw std::invalid_argument("unsupported ADetailer setting: " + item.key()); + } + const auto& configured = runtime.default_gen_params->ad_model_path; + if (configured.empty()) + throw std::invalid_argument("ADetailer requires a server-configured --ad-model"); + const auto path = fs::u8path(configured); + if (model.find_first_of("/\\") != std::string::npos || + (lower_ascii(model) != lower_ascii(path.filename().u8string()) && lower_ascii(model) != lower_ascii(path.stem().u8string() + ".pt"))) { + throw std::invalid_argument("ADetailer model must match the server-configured detector"); + } + if (!text("ad_model_classes", "").empty()) + throw std::invalid_argument("ad_model_classes is not supported"); + if (boolean("ad_use_checkpoint", false)) { + const auto checkpoint = config.contains("ad_checkpoint") && config.at("ad_checkpoint").is_null() ? "" : text("ad_checkpoint", "Use same checkpoint"); + if (!checkpoint.empty() && checkpoint != "None" && checkpoint != "Use same checkpoint") + throw std::invalid_argument("ADetailer checkpoint switching is not supported"); + } + if (boolean("ad_use_vae", false)) { + const auto vae = config.contains("ad_vae") && config.at("ad_vae").is_null() ? "" : text("ad_vae", "Use same VAE"); + if (!vae.empty() && vae != "None" && vae != "Use same VAE") + throw std::invalid_argument("ADetailer VAE switching is not supported"); + } + if (boolean("ad_use_noise_multiplier", false) && number("ad_noise_multiplier", 1, 0.5, 1.5) != 1) + throw std::invalid_argument("ADetailer noise multiplier is not supported"); + if (boolean("ad_restore_face", false)) + throw std::invalid_argument("ADetailer face restoration is not supported"); + if (text("ad_controlnet_model", "None") != "None") + throw std::invalid_argument("ADetailer ControlNet is not supported"); + + ServerADetailerPass pass; + pass.prompt = text("ad_prompt", ""); + pass.negative_prompt = text("ad_negative_prompt", ""); + pass.width = request.gen_params.get_resolved_width(); + pass.height = request.gen_params.get_resolved_height(); + pass.strength = static_cast(number("ad_denoising_strength", 0.4, 0, 1)); + auto append = [&](const char* key, const std::string& value) { + if (!pass.extra_args.empty()) + pass.extra_args += ","; + pass.extra_args += std::string(key) + "=" + value; + }; + auto append_number = [&](const char* target, const char* key, double fallback, double minimum, double maximum) { + append(target, json(number(key, fallback, minimum, maximum)).dump()); + }; + auto append_integer = [&](const char* target, const char* key, int fallback, int minimum, int maximum = std::numeric_limits::max()) { + append(target, std::to_string(integer(key, fallback, minimum, maximum))); + }; + append_number("confidence", "ad_confidence", 0.3, 0, 1); + append_integer("mask_k_largest", "ad_mask_k", 0, 0); + const auto filter = text("ad_mask_filter_method", "Area"); + if (filter != "Area" && filter != "Confidence") + throw std::invalid_argument("invalid ad_mask_filter_method"); + append("mask_filter_method", lower_ascii(filter)); + const auto min_ratio = number("ad_mask_min_ratio", 0, 0, 1); + const auto max_ratio = number("ad_mask_max_ratio", 1, 0, 1); + if (min_ratio > max_ratio) + throw std::invalid_argument("ad_mask_min_ratio exceeds ad_mask_max_ratio"); + append("mask_min_ratio", json(min_ratio).dump()); + append("mask_max_ratio", json(max_ratio).dump()); + // Bound pixel operations before converting them to the native mask parameters. + const int mask_limit = std::clamp(std::max(pass.width, pass.height), 32, 16384); + append_integer("dilate_erode", "ad_dilate_erode", 4, -mask_limit, mask_limit); + append_integer("x_offset", "ad_x_offset", 0, -mask_limit, mask_limit); + append_integer("y_offset", "ad_y_offset", 0, -mask_limit, mask_limit); + append_integer("mask_blur", "ad_mask_blur", 4, 0, mask_limit); + append_integer("inpaint_padding", "ad_inpaint_only_masked_padding", 32, 0, mask_limit); + append("inpaint_only_masked", boolean("ad_inpaint_only_masked", true) ? "true" : "false"); + const auto merge = text("ad_mask_merge_invert", "None"); + if (merge != "None" && merge != "Merge" && merge != "Merge and Invert") + throw std::invalid_argument("invalid ad_mask_merge_invert"); + append("mask_mode", merge == "None" ? "none" : merge == "Merge" ? "merge" + : "merge_invert"); + if (boolean("ad_use_inpaint_width_height", false)) { + pass.width = integer("ad_inpaint_width", 512, 1, k_max_generation_dimension); + pass.height = integer("ad_inpaint_height", 512, 1, k_max_generation_dimension); + } + if (boolean("ad_use_steps", false)) + append_integer("steps", "ad_steps", 28, 1, k_max_adetailer_steps); + if (boolean("ad_use_cfg_scale", false)) + append_number("cfg_scale", "ad_cfg_scale", 7, 0, std::numeric_limits::max()); + if (boolean("ad_use_clip_skip", false)) { + pass.clip_skip = integer("ad_clip_skip", 1, 1); + if (pass.clip_skip > 12) + throw std::invalid_argument("ad_clip_skip is out of range"); + } + if (boolean("ad_use_sampler", false)) { + std::string sampler = text("ad_sampler", "DPM++ 2M Karras"); + if (sampler != "Use same sampler") { + const std::string suffix = " karras"; + const bool legacy_karras = sampler.size() > suffix.size() && lower_ascii(sampler.substr(sampler.size() - suffix.size())) == suffix; + if (legacy_karras) + sampler.resize(sampler.size() - suffix.size()); + const auto method = get_sdapi_sample_method(sampler); + if (method == SAMPLE_METHOD_COUNT) + throw std::invalid_argument("unsupported ADetailer sampler"); + append("sample_method", sd_sample_method_name(method)); + if (legacy_karras && text("ad_scheduler", "Use same scheduler") == "Use same scheduler") + append("scheduler", sd_scheduler_name(KARRAS_SCHEDULER)); + } + const auto scheduler = text("ad_scheduler", "Use same scheduler"); + if (scheduler == "Automatic") { + pass.automatic_scheduler = true; + } else if (scheduler != "Use same scheduler") { + const auto value = str_to_scheduler(scheduler.c_str()); + if (value == SCHEDULER_COUNT) + throw std::invalid_argument("unsupported ADetailer scheduler"); + append("scheduler", sd_scheduler_name(value)); + } + } + request.adetailer_passes.push_back(std::move(pass)); + } + if (!request.adetailer_passes.empty() && img2img && args.size() > 1 && args[1].is_boolean() && args[1].get()) { + throw std::invalid_argument("ADetailer skip img2img is not supported"); + } + request.adetailer_enabled = !request.adetailer_passes.empty(); + return true; + } catch (const std::exception& error) { + error_message = std::string("invalid ADetailer args: ") + error.what(); + return false; + } +} + static bool build_sdapi_img_gen_request(const json& j, ServerRuntime& runtime, bool img2img, @@ -259,49 +458,31 @@ static bool build_sdapi_img_gen_request(const json& j, return false; } - return true; + return parse_sdapi_adetailer(j, runtime, img2img, request, error_message) && + validate_server_img_gen_request(runtime, request, error_message); } -static nlohmann::json prepare_info_field(const SDContextParams& ctx_params, - const SDGenerationParams& gen_params, - bool img2img) { - nlohmann::json jsoninfo = nlohmann::json::object(); - jsoninfo["prompt"] = gen_params.prompt; - if (!gen_params.negative_prompt.empty()) { - jsoninfo["negative_prompt"] = gen_params.negative_prompt; - } - jsoninfo["seed"] = gen_params.seed; - jsoninfo["cfg_scale"] = gen_params.sample_params.guidance.txt_cfg; - jsoninfo["width"] = gen_params.get_resolved_width(); - jsoninfo["height"] = gen_params.get_resolved_height(); - jsoninfo["steps"] = gen_params.sample_params.sample_steps; - jsoninfo["sampler_name"] = sd_sample_method_name(gen_params.sample_params.sample_method); - if (gen_params.clip_skip != -1) { - jsoninfo["clip_skip"] = gen_params.clip_skip; - } - if (gen_params.sample_params.scheduler != scheduler_t::SCHEDULER_COUNT) { - jsoninfo["extra_generation_params"] = nlohmann::json::object(); - jsoninfo["extra_generation_params"]["Schedule type"] = sd_scheduler_name(gen_params.sample_params.scheduler); - } - if (img2img) { - jsoninfo["denoising_strength"] = gen_params.strength; - } - // not clear what should happen if we have both model and diffusion_model - if (!ctx_params.diffusion_model_path.empty()) { - jsoninfo["sd_model_name"] = sd_basename(ctx_params.diffusion_model_path); - } else if (!ctx_params.model_path.empty()) { - jsoninfo["sd_model_name"] = sd_basename(ctx_params.model_path); - } - if (!ctx_params.vae_path.empty()) { - jsoninfo["sd_vae_name"] = sd_basename(ctx_params.vae_path); - } - jsoninfo["version"] = "stable-diffusion.cpp"; - - jsoninfo["infotexts"] = nlohmann::json::array(); - jsoninfo["all_prompts"] = nlohmann::json::array(); - jsoninfo["all_negative_prompts"] = nlohmann::json::array(); - jsoninfo["all_seeds"] = nlohmann::json::array(); - return jsoninfo; +static nlohmann::json prepare_info_field(const std::string& generation_info) { + const auto info = nlohmann::json::parse(generation_info); + const auto& sampling = info.at("sample_params"); + nlohmann::json result = {{"prompt", info.at("prompt")}, + {"negative_prompt", info.at("negative_prompt")}, + {"seed", generation_seed(generation_info)}, + {"cfg_scale", sampling.at("guidance").at("txt_cfg")}, + {"width", info.at("width")}, + {"height", info.at("height")}, + {"steps", sampling.contains("custom_sigmas") ? json(sampling.at("custom_sigmas").size() - 1) : sampling.at("sample_steps")}, + {"sampler_name", sampling.at("sample_method")}, + {"version", "stable-diffusion.cpp"}}; + if (info.contains("clip_skip")) + result["clip_skip"] = info.at("clip_skip"); + if (info.contains("strength")) + result["denoising_strength"] = info.at("strength"); + if (sampling.contains("scheduler")) + result["extra_generation_params"]["Schedule type"] = sampling.at("scheduler"); + for (const auto* key : {"infotexts", "all_prompts", "all_negative_prompts", "all_seeds", "generation_info"}) + result[key] = nlohmann::json::array(); + return result; } struct SDAPIPreviewState { @@ -343,13 +524,13 @@ static void preview_callback(int step, int frame_count, sd_image_t* frames, bool std::lock_guard lock(g_preview.mutex); g_preview.b64 = std::move(b64); g_preview.step = static_cast(std::min(step < 0 ? -static_cast(step) : step, - std::max(0, info.total_steps))); + std::max(0, info.total_steps))); g_preview.total = std::max(0, info.total_steps); g_preview.job_no = std::max(0, info.sample_pass - 1); } struct SDAPIPreviewGuard { - SDAPIPreviewGuard(const std::string& id_task, const sd_img_gen_params_t& params, const std::string& job_type) { + SDAPIPreviewGuard(const std::string& id_task, const sd_img_gen_params_t& params, const std::string& job_type, bool detailing = false) { std::lock_guard lock(g_preview.mutex); g_preview.id_task = id_task; g_preview.total = 0; @@ -358,11 +539,20 @@ struct SDAPIPreviewGuard { g_preview.step = 0; g_preview.active = true; // Each batch item has a base pass and, when enabled, a hires pass. - g_preview.job_count = static_cast(std::max(1, params.batch_count)) * (params.hires.enabled ? 2 : 1); + g_preview.job_count = static_cast(std::max(1, params.batch_count)) * (params.hires.enabled ? 2 : 1) + (detailing ? 1 : 0); g_preview.job_no = 0; sd_set_preview_callback(preview_callback, PREVIEW_PROJ, 1, true, false, nullptr); } + void begin_detailing() { + sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr); + std::lock_guard lock(g_preview.mutex); + g_preview.job_type = "adetailer"; + g_preview.job_no = static_cast(g_preview.job_count - 1); + g_preview.step = 0; + g_preview.total = 0; + } + ~SDAPIPreviewGuard() { sd_set_preview_callback(nullptr, PREVIEW_NONE, 1, false, false, nullptr); std::lock_guard lock(g_preview.mutex); @@ -410,13 +600,19 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { { std::lock_guard lock(*runtime->sd_ctx_mutex); - SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img"); - sd_image_t* raw_results = nullptr; - if (!generate_image(runtime->sd_ctx, &img_gen_params, &raw_results, &num_results)) { - raw_results = nullptr; - num_results = 0; + const bool detailing = request.adetailer_enabled && !runtime->default_gen_params->ad_model_path.empty(); + SDAPIPreviewGuard preview_guard(id_task, img_gen_params, img2img ? "img2img" : "txt2img", detailing); + sd_image_generation_result_t* generated = nullptr; + generate_image(runtime->sd_ctx, &img_gen_params, &generated); + results.adopt(generated, runtime->ctx_params->lora_model_dir); + num_results = results.count(); + if (detailing && !results.empty()) + preview_guard.begin_detailing(); + if (!results.empty() && !apply_server_adetailer(*runtime, request, img_gen_params, results)) { + res.status = 500; + res.set_content(R"({"error":"ADetailer failed"})", "application/json"); + return; } - results.adopt(raw_results, num_results); } if (results.empty()) { @@ -428,9 +624,8 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { json out; out["images"] = json::array(); out["parameters"] = j; - json jsoninfo = prepare_info_field(*runtime->ctx_params, request.gen_params, img2img); + json jsoninfo = prepare_info_field(results.generation_info(0)); - int images_per_batch = request.gen_params.batch_count > 0 ? std::max(1, num_results / request.gen_params.batch_count) : 1; for (int i = 0; i < num_results; ++i) { if (results[i].data == nullptr) { continue; @@ -438,16 +633,14 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { bool embed_meta = request.gen_params.embed_image_metadata; - std::string params = get_image_params(*runtime->ctx_params, - request.gen_params, - request.gen_params.seed + i / images_per_batch); + std::string params = generation_info_to_webui(results.generation_info(i)); auto image_bytes = encode_image_to_vector(EncodedImageFormat::PNG, results[i].data, results[i].width, results[i].height, results[i].channel, - embed_meta ? params : ""); + embed_meta ? results.generation_info(i) : ""); if (image_bytes.empty()) { LOG_ERROR("write image to mem failed"); @@ -456,10 +649,12 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) { out["images"].push_back(base64_encode(image_bytes)); - jsoninfo["infotexts"][i] = params; - jsoninfo["all_seeds"][i] = request.gen_params.seed + i; - jsoninfo["all_prompts"][i] = request.gen_params.prompt; - jsoninfo["all_negative_prompts"][i] = request.gen_params.negative_prompt; + const auto actual = json::parse(results.generation_info(i)); + jsoninfo["infotexts"].push_back(params); + jsoninfo["all_seeds"].push_back(generation_seed(results.generation_info(i))); + jsoninfo["generation_info"].push_back(actual); + jsoninfo["all_prompts"].push_back(actual["prompt"]); + jsoninfo["all_negative_prompts"].push_back(actual["negative_prompt"]); } // not a mistake: it is supposed to be a string in json format diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index 7e93e5edbf..0dda9a95be 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -178,6 +178,7 @@ static json make_vid_gen_defaults_json(const SDGenerationParams& defaults, const {"seed", defaults.seed}, {"video_frames", defaults.video_frames}, {"image_preprocess", defaults.image_preprocess}, + {"ref_image_args", defaults.ref_image_args}, {"fps", defaults.fps}, {"moe_boundary", defaults.moe_boundary}, {"vace_strength", defaults.vace_strength}, @@ -206,22 +207,26 @@ static json make_img_gen_features_json() { {"hires", true}, {"cache", true}, {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, {"preview", true}, }; } +/** Builds the default feature set advertised for vid_gen requests. */ static json make_vid_gen_features_json() { return { {"init_image", true}, {"end_image", true}, {"control_frames", true}, + {"ref_videos", true}, + {"ref_audios", true}, + {"ref_video_audios", true}, {"high_noise_sample_params", true}, {"lora", true}, {"vae_tiling", true}, {"cache", true}, {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_generating", true}, {"preview", true}, }; } @@ -234,6 +239,7 @@ static json make_preview_modes_json() { return modes; } +/** Builds the /capabilities payload: model, defaults, limits and feature maps. */ static json make_capabilities_json(ServerRuntime& runtime) { refresh_lora_cache(runtime); refresh_upscaler_cache(runtime); @@ -350,8 +356,10 @@ static json make_capabilities_json(ServerRuntime& runtime) { const json generation_features = model_generation_features(runtime); auto recipe_features = [&](json features, bool video) { features["init_image"] = generation_features.at(video ? "init_image" : "img2img"); - for (const char* key : {"mask_image", "ref_images", "end_image", "control_frames", "lora", "vae_tiling", "cache", "preview"}) { - if (features.contains(key) || std::string(key) == "ref_images") + for (const char* key : {"mask_image", "ref_images", "end_image", "control_frames", "ref_videos", "ref_audios", "ref_video_audios", "lora", "vae_tiling", "cache", "preview"}) { + if (features.contains(key) || std::string(key) == "ref_images" || + (video && (std::string(key) == "ref_videos" || std::string(key) == "ref_audios" || + std::string(key) == "ref_video_audios"))) features[key] = generation_features.at(key); } if (video) @@ -380,8 +388,8 @@ static json make_capabilities_json(ServerRuntime& runtime) { json top_level_defaults = json::object(); json top_level_output_formats = json::array(); json top_level_features = { - {"cancel_queued", true}, - {"cancel_generating", false}, + {"cancel_queued", true}, + {"cancel_generating", true}, }; std::string current_mode = ""; if (supports_img) { @@ -409,14 +417,14 @@ static json make_capabilities_json(ServerRuntime& runtime) { result["defaults"] = top_level_defaults; result["defaults_by_mode"] = defaults_by_mode; result["limits"] = { - {"min_width", 64}, - {"max_width", 4096}, - {"min_height", 64}, - {"max_height", 4096}, - {"max_batch_count", 8}, - {"max_queue_size", manager.max_pending_jobs}, - {"max_upscale_width", k_max_upscale_dimension}, - {"max_upscale_height", k_max_upscale_dimension}, + {"min_width", 64}, + {"max_width", k_max_generation_dimension}, + {"min_height", 64}, + {"max_height", k_max_generation_dimension}, + {"max_batch_count", 8}, + {"max_queue_size", manager.max_pending_jobs}, + {"max_upscale_width", k_max_upscale_dimension}, + {"max_upscale_height", k_max_upscale_dimension}, }; result["samplers"] = samplers; result["schedulers"] = schedulers; @@ -457,7 +465,7 @@ static bool parse_img_gen_request(const json& body, error_message = "invalid generation parameters"; return false; } - return true; + return validate_server_img_gen_request(runtime, request, error_message); } static bool parse_vid_gen_request(const json& body, @@ -519,7 +527,7 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { const bool valid = value.is_number_unsigned() ? value.get() <= static_cast(std::numeric_limits::max()) : value.is_number_integer() && value.get() >= std::numeric_limits::min() && - value.get() <= std::numeric_limits::max(); + value.get() <= std::numeric_limits::max(); if (!valid) { res.status = 400; res.set_content(json({{"error", std::string(key) + " must be a 32-bit integer"}}).dump(), "application/json"); @@ -608,19 +616,20 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } SDImageOwner current(input.release()); + std::string generation_info; for (int i = 0; i < repeats; ++i) { - sd_image_t* out_images = nullptr; - int out_count = 0; - if (!upscale(upscaler_ctx.get(), current.get(), (uint32_t)factor, &out_images, &out_count) || - out_count <= 0 || out_images[0].data == nullptr) { - free_sd_images(out_images, out_count); + sd_image_generation_result_t* generated = nullptr; + if (!upscale(upscaler_ctx.get(), current.get(), (uint32_t)factor, &generated) || + !generated || generated->count <= 0 || generated->items[0].image.data == nullptr) { + free_sd_image_generation_result(generated); res.status = 500; res.set_content(R"({"error":"upscale failed"})", "application/json"); return; } - sd_image_t produced = out_images[0]; - out_images[0] = {0, 0, 0, nullptr}; - free_sd_images(out_images, out_count); + generation_info = merge_generation_info(generation_info, generated->items[0].generation_info); + sd_image_t produced = generated->items[0].image; + generated->items[0].image = {}; + free_sd_image_generation_result(generated); current.reset(produced); } @@ -630,12 +639,12 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { auto image_bytes = encode_image_to_vector(format == "jpeg" ? EncodedImageFormat::JPEG : format == "webp" ? EncodedImageFormat::WEBP : EncodedImageFormat::PNG, - result.data, - result.width, - result.height, - result.channel, - "", - compression); + result.data, + result.width, + result.height, + result.channel, + generation_info, + compression); if (image_bytes.empty()) { res.status = 500; res.set_content(R"({"error":"the result could not be encoded"})", "application/json"); @@ -643,13 +652,14 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } json out; - out["upscaler"] = used_name; - out["scale"] = factor; - out["repeats"] = repeats; - out["width"] = result.width; - out["height"] = result.height; - out["output_format"] = format; - json images = json::array(); + out["generation_info"] = json::parse(generation_info); + out["upscaler"] = used_name; + out["scale"] = factor; + out["repeats"] = repeats; + out["width"] = result.width; + out["height"] = result.height; + out["output_format"] = format; + json images = json::array(); images.push_back({{"index", 0}, {"b64_json", base64_encode(image_bytes)}}); out["images"] = std::move(images); res.set_content(out.dump(), "application/json"); @@ -840,8 +850,9 @@ void register_sdcpp_api_endpoints(httplib::Server& svr, ServerRuntime& rt) { } if (job.status == AsyncJobStatus::Generating) { - res.status = 409; - res.set_content(R"({"error":"job is currently generating and cannot be interrupted yet"})", "application/json"); + job.cancel_requested.store(true); + res.status = 200; + res.set_content(make_async_job_json(manager, job).dump(), "application/json"); return; } diff --git a/examples/server/runtime.cpp b/examples/server/runtime.cpp index 34140ea353..b187ff85a3 100644 --- a/examples/server/runtime.cpp +++ b/examples/server/runtime.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -18,6 +19,156 @@ namespace fs = std::filesystem; +bool validate_server_img_gen_request(const ServerRuntime& runtime, + const ImgGenJobRequest& request, + std::string& error_message) { + auto dimensions_valid = [&](double width, double height) { + if (!std::isfinite(width) || !std::isfinite(height) || + width < 1 || height < 1 || width > k_max_generation_dimension || height > k_max_generation_dimension) { + error_message = "generation width and height must be between 1 and " + std::to_string(k_max_generation_dimension); + return false; + } + return true; + }; + const auto& params = request.gen_params; + const int width = params.get_resolved_width(); + const int height = params.get_resolved_height(); + if (!dimensions_valid(width, height)) { + return false; + } + if (params.hires_enabled) { + const double hires_width = params.hires_width > 0 ? params.hires_width : params.hires_height > 0 ? params.hires_height + : std::round(width * params.hires_scale); + const double hires_height = params.hires_height > 0 ? params.hires_height : params.hires_width > 0 ? params.hires_width + : std::round(height * params.hires_scale); + if (!dimensions_valid(hires_width, hires_height)) { + return false; + } + } + if (!request.adetailer_enabled || runtime.default_gen_params->ad_model_path.empty()) { + return true; + } + auto pass_valid = [&](const ServerADetailerPass& pass) { + if (!dimensions_valid(pass.width, pass.height)) { + return false; + } + // Match the native key normalization, but leave full option parsing to ADetailer. + auto trim = [](const std::string& text) { + const auto first = text.find_first_not_of(" \t\r\n"); + return first == std::string::npos ? std::string() : text.substr(first, text.find_last_not_of(" \t\r\n") - first + 1); + }; + std::stringstream stream(pass.extra_args); + std::string item; + while (std::getline(stream, item, ',')) { + const auto equals = item.find('='); + if (equals == std::string::npos) { + continue; + } + auto key = trim(item.substr(0, equals)); + std::transform(key.begin(), key.end(), key.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); }); + if (key != "inpaint_width" && key != "inpaint_height" && key != "steps") { + continue; + } + const auto value = trim(item.substr(equals + 1)); + const int minimum = key == "steps" ? 0 : 1; + const int maximum = key == "steps" ? k_max_adetailer_steps : k_max_generation_dimension; + try { + size_t used = 0; + const int parsed = std::stoi(value, &used); + if (used == value.size() && parsed >= minimum && parsed <= maximum) { + continue; + } + } catch (const std::exception&) { + } + error_message = "ADetailer " + key + " must be between " + std::to_string(minimum) + " and " + std::to_string(maximum); + return false; + } + return true; + }; + if (request.adetailer_passes.empty()) { + ServerADetailerPass pass; + pass.extra_args = params.extra_ad_args; + return pass_valid(pass); + } + return std::all_of(request.adetailer_passes.begin(), request.adetailer_passes.end(), pass_valid); +} + +bool apply_server_adetailer(ServerRuntime& runtime, + const ImgGenJobRequest& request, + const sd_img_gen_params_t& params, + SDImageVec& results) { + const auto& gen_params = request.gen_params; + const auto& ad_model_path = runtime.default_gen_params->ad_model_path; + if (!request.adetailer_enabled || ad_model_path.empty()) { + return true; + } + std::string error_message; + if (!validate_server_img_gen_request(runtime, request, error_message)) { + LOG_ERROR("%s", error_message.c_str()); + return false; + } + + std::vector passes = request.adetailer_passes; + if (passes.empty()) { + ServerADetailerPass pass; + pass.prompt = gen_params.ad_prompt; + pass.negative_prompt = gen_params.ad_negative_prompt; + pass.extra_args = gen_params.extra_ad_args; + passes.push_back(std::move(pass)); + } + + ADetailerCtxPtr ad_ctx(new_adetailer_ctx(ad_model_path.c_str(), + runtime.ctx_params->n_threads, + runtime.ctx_params->effective_backend.c_str(), + runtime.ctx_params->effective_params_backend.c_str())); + if (ad_ctx == nullptr) { + LOG_ERROR("new_adetailer_ctx failed"); + return false; + } + + for (const auto& pass : passes) { + sd_adetailer_params_t ad_params{}; + ad_params.prompt = pass.prompt.empty() ? nullptr : pass.prompt.c_str(); + ad_params.negative_prompt = pass.negative_prompt.empty() ? nullptr : pass.negative_prompt.c_str(); + ad_params.extra_ad_args = pass.extra_args.c_str(); + for (int i = 0; i < results.count(); ++i) { + if (results[i].data == nullptr) { + continue; + } + sd_img_gen_params_t ad_generation_params = params; + ad_generation_params.seed = generation_seed(results.generation_info(i)); + ad_generation_params.width = pass.width; + ad_generation_params.height = pass.height; + ad_generation_params.strength = pass.strength; + if (pass.clip_skip > 0) { + ad_generation_params.clip_skip = pass.clip_skip; + } + if (pass.automatic_scheduler) { + ad_generation_params.sample_params.scheduler = SCHEDULER_COUNT; + } + + sd_image_generation_result_t* detailed = nullptr; + if (!adetail_image(ad_ctx.get(), + runtime.sd_ctx, + results[i], + &ad_params, + &ad_generation_params, + &detailed) || + detailed == nullptr || detailed->count <= 0 || detailed->items[0].image.data == nullptr) { + free_sd_image_generation_result(detailed); + LOG_ERROR("ADetailer failed for image %d", i + 1); + return false; + } + results.append_generation_info(i, detailed->items[0].generation_info, runtime.ctx_params->lora_model_dir); + free(results[i].data); + results[i] = detailed->items[0].image; + detailed->items[0].image = {}; + free_sd_image_generation_result(detailed); + } + } + return true; +} + static std::string lower_ascii(std::string value) { std::transform(value.begin(), value.end(), value.begin(), [](unsigned char c) { return static_cast(std::tolower(c)); diff --git a/examples/server/runtime.h b/examples/server/runtime.h index d9272f6c24..8986f38270 100644 --- a/examples/server/runtime.h +++ b/examples/server/runtime.h @@ -67,18 +67,38 @@ struct ServerRuntime { ServerModels* models; }; +struct ServerADetailerPass { + std::string prompt; + std::string negative_prompt; + std::string extra_args; + int width = 512; + int height = 512; + float strength = 0.4f; + int clip_skip = -1; + bool automatic_scheduler = false; +}; + struct ImgGenJobRequest { SDGenerationParams gen_params; std::string output_format = "png"; int output_compression = 100; std::string preview_mode = "none"; int preview_interval = 1; + bool adetailer_enabled = true; + std::vector adetailer_passes; sd_img_gen_params_t to_sd_img_gen_params_t() { return gen_params.to_sd_img_gen_params_t(); } }; +constexpr int k_max_generation_dimension = 4096; +constexpr int k_max_adetailer_steps = 100; + +bool validate_server_img_gen_request(const ServerRuntime& runtime, + const ImgGenJobRequest& request, + std::string& error_message); + struct VidGenJobRequest { SDGenerationParams gen_params; std::string output_format = "webm"; @@ -111,3 +131,7 @@ void refresh_lora_cache(ServerRuntime& rt); std::string get_lora_full_path(ServerRuntime& rt, const std::string& path); void refresh_upscaler_cache(ServerRuntime& rt); int64_t unix_timestamp_now_ms(); +bool apply_server_adetailer(ServerRuntime& runtime, + const ImgGenJobRequest& request, + const sd_img_gen_params_t& params, + SDImageVec& results); diff --git a/ggml b/ggml index dc998771e7..135ce1e0e3 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit dc998771e700d49bbdf9a39cede22bb5b3311767 +Subproject commit 135ce1e0e3f2b0693d3de93d4f6b5fa0a1947820 diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 404d8cbddc..f4ea0eaf13 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -94,6 +94,9 @@ enum prediction_t { SEFI_FLOW_PRED, MINIT2I_FLOW_PRED, SENSENOVA_U1_FLOW_PRED, + FLUX3_FLOW_PRED, + H3_AV_FLOW_PRED, + MING_IMAGE_FLOW_PRED, PREDICTION_COUNT }; @@ -269,6 +272,25 @@ typedef struct { uint8_t* data; } sd_image_t; +typedef struct { + sd_image_t image; + // Owned UTF-8 JSON using the sdcpp.generation/v1 schema. + char* generation_info; +} sd_image_result_t; + +typedef struct { + sd_image_result_t* items; + int count; +} sd_image_generation_result_t; + +typedef struct { + sd_image_t* frames; + int frame_count; + sd_audio_t* audio; + int fps; + char* generation_info; +} sd_video_generation_result_t; + typedef struct { // Semicolon-separated target=...,key=value rules. NULL preserves defaults. const char* rules; @@ -456,6 +478,7 @@ typedef struct { bool circular_x; bool circular_y; sd_image_preprocess_params_t image_preprocess; + const char* ref_image_args; } sd_vid_gen_params_t; typedef struct sd_ctx_t sd_ctx_t; @@ -486,8 +509,6 @@ SD_API int32_t sd_get_num_physical_cores(); SD_API const char* sd_get_system_info(); SD_API bool sd_ctx_supports_image_generation(const sd_ctx_t* sd_ctx); SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx); -/** Returns whether the loaded model consumes ref_images as conditioning inputs. */ -SD_API bool sd_ctx_supports_ref_image_conditioning(const sd_ctx_t* sd_ctx); // ControlNet hot-swap APIs are not safe to call while generation is in flight. SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path); @@ -534,8 +555,12 @@ SD_API void sd_img_gen_params_init(sd_img_gen_params_t* sd_img_gen_params); SD_API char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params); SD_API bool generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* sd_img_gen_params, - sd_image_t** images_out, - int* num_images_out); + sd_image_generation_result_t** result_out); + +// Results own their pixels, audio and generation_info, independently of the context. +// Generation calls require result_out and set it to NULL on failure. +SD_API void free_sd_image_generation_result(sd_image_generation_result_t* result); +SD_API void free_sd_video_generation_result(sd_video_generation_result_t* result); enum sd_cancel_mode_t { // Stop the current generation as soon as possible. @@ -549,13 +574,9 @@ enum sd_cancel_mode_t { SD_API void sd_cancel_generation(sd_ctx_t* sd_ctx, enum sd_cancel_mode_t mode); SD_API void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params); -// If non-NULL, fps_out receives the effective encoding frame rate before preview callbacks. SD_API bool generate_video(sd_ctx_t* sd_ctx, const sd_vid_gen_params_t* sd_vid_gen_params, - sd_image_t** frames_out, - int* num_frames_out, - sd_audio_t** audio_out, - int* fps_out); + sd_video_generation_result_t** result_out); typedef struct upscaler_ctx_t upscaler_ctx_t; @@ -570,8 +591,7 @@ SD_API void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx); SD_API bool upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor, - sd_image_t** images_out, - int* num_images_out); + sd_image_generation_result_t** result_out); 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. @@ -595,8 +615,7 @@ SD_API bool adetail_image(adetailer_ctx_t* adetailer_ctx, sd_image_t input_image, const sd_adetailer_params_t* adetailer_params, const sd_img_gen_params_t* inpaint_params, - sd_image_t** images_out, - int* num_images_out); + sd_image_generation_result_t** result_out); SD_API bool convert(const char* input_path, const char* vae_path, @@ -640,6 +659,11 @@ SD_API const char* sd_version(void); // terminator. Passing nullptr or buffer_size 0 only queries the required size. SD_API size_t sd_list_devices(char* buffer, size_t buffer_size); +// Register comma-separated RPC endpoints (host:port) before listing devices. +// Repeated endpoints are safe; nullptr or an empty string is a no-op. +// Returns false if RPC is unavailable or an endpoint has no usable devices. +SD_API bool sd_add_rpc_servers(const char* servers); + // for C API, caller needs to call free_sd_images to free the memory after use // This helps avoid CRT problems on Windows when memory is allocated in the library but freed in the caller, which may use a different CRT. SD_API void free_sd_images(sd_image_t* result_images, int num_images); diff --git a/scripts/ci-cuda-cache.sh b/scripts/ci-cuda-cache.sh new file mode 100644 index 0000000000..6dda05525c --- /dev/null +++ b/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 diff --git a/scripts/format-code.ps1 b/scripts/format-code.ps1 index 91c5f3b03a..479433b3f1 100644 --- a/scripts/format-code.ps1 +++ b/scripts/format-code.ps1 @@ -8,8 +8,17 @@ $patterns = @( "src/core/*.cpp" "src/core/*.h" "src/core/*.hpp" + "src/core/cpu/*.cpp" + "src/core/cpu/*.h" + "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 +32,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 +53,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") { diff --git a/scripts/format-code.sh b/scripts/format-code.sh index 568a0e54bb..29b3313a9d 100644 --- a/scripts/format-code.sh +++ b/scripts/format-code.sh @@ -8,16 +8,19 @@ 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/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" diff --git a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch index 190aac4918..0264e57339 100644 --- a/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch +++ b/scripts/unsloth/ggml-patches/0002-ggml-cuda-fa-longseq.patch @@ -33,11 +33,15 @@ tiles, several Q tiles and GQA broadcast. 4 files changed, 184 insertions(+), 15 deletions(-) create mode 100644 src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu -[Ported onto leejet/ggml dc998771 (stable-diffusion.cpp master-982): the diff below is - the whole patch as one diff, regenerated from the port.] +[Ported onto leejet/ggml 135ce1e0 (stable-diffusion.cpp master-1050): the diff below is + the whole patch as one diff, regenerated from the port. + Upstream ggml 0d84d52a (wide flash-attention tiles) added the DKQ=128 ncols1=128 + config row, its instance (fattn-mma-wide.cu) and ggml_cuda_fattn_mma_shared_size; the + port reuses them (shared_size gains an nstages override for the long-sequence variant) + and the instance file here now only adds DKQ=64 ncols1=128.] --- diff --git a/src/ggml-cuda/fattn-common.cuh b/src/ggml-cuda/fattn-common.cuh -index 41d94833..82f43137 100644 +index 8bfe7e30..1840b69c 100644 --- a/src/ggml-cuda/fattn-common.cuh +++ b/src/ggml-cuda/fattn-common.cuh @@ -981,7 +981,10 @@ template @@ -63,21 +67,21 @@ index 41d94833..82f43137 100644 int parallel_blocks = max_blocks_per_sm; diff --git a/src/ggml-cuda/fattn-mma-f16.cuh b/src/ggml-cuda/fattn-mma-f16.cuh -index 296aa2b3..e011c8d5 100644 +index 391c75be..ca372be9 100644 --- a/src/ggml-cuda/fattn-mma-f16.cuh +++ b/src/ggml-cuda/fattn-mma-f16.cuh -@@ -61,6 +61,10 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co +@@ -62,6 +62,10 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 32, 128, 2, 64, 64, 64, 64, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 64, 128, 2, 64, 64, 64, 64, 2, true); -+ // ncols == 128 is only used by the long-sequence variant (ncols2 == 1, no mask), see ggml_cuda_fattn_mma_use_longseq: ++ // ncols == 128 for DKQ == 64 is only used by the long-sequence variant (ncols2 == 1, no mask), see ++ // ggml_cuda_fattn_mma_use_longseq; the DKQ == 128 row is the wide tile at the top of this table. + GGML_CUDA_FATTN_MMA_CONFIG_CASE( 64, 64, 128, 256, 1, 64, 32, 32, 32, 2, true); -+ GGML_CUDA_FATTN_MMA_CONFIG_CASE(128, 128, 128, 256, 1, 64, 64, 64, 64, 2, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 8, 64, 4, 64, 96, 64, 64, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 16, 64, 4, 32, 96, 64, 64, 2, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(192, 128, 32, 128, 2, 32, 96, 64, 64, 2, true); -@@ -378,6 +382,24 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages( +@@ -379,6 +383,24 @@ static constexpr __device__ int ggml_cuda_fattn_mma_get_nstages( #endif // CP_ASYNC_AVAILABLE } @@ -102,7 +106,7 @@ index 296aa2b3..e011c8d5 100644 // ------------------------------------------------------------------------------------------------------------------ template -@@ -573,7 +595,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( +@@ -574,7 +596,7 @@ static __device__ __forceinline__ void flash_attn_ext_f16_load_mask( template struct mma_tile_sizes { +@@ -1187,7 +1209,8 @@ template struct mma_tile_sizes { }; #endif // defined(TURING_MMA_AVAILABLE) @@ -130,7 +134,7 @@ index 296aa2b3..e011c8d5 100644 static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( const float2 * const __restrict__ Q_f2, const half2 * const __restrict__ K_h2, -@@ -1232,7 +1255,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( +@@ -1233,7 +1256,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols); constexpr int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols); constexpr bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols); @@ -141,7 +145,7 @@ index 296aa2b3..e011c8d5 100644 if (cols_per_warp > ncols) { NO_DEVICE_CODE; -@@ -1331,6 +1356,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( +@@ -1332,6 +1357,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( int kb0 = kb0_start; @@ -151,7 +155,7 @@ index 296aa2b3..e011c8d5 100644 // Preload mask and K data for first iteration when using cp_async with multiple stages: if constexpr (nstages > 1) { static_assert(!use_sparse, "sparse gather not implemented for multi-stage loading"); -@@ -1338,16 +1366,54 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( +@@ -1339,16 +1367,54 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( constexpr bool use_cp_async = true; constexpr bool oob_check = false; constexpr int k_VKQ_sup = nbatch_fa; @@ -212,7 +216,7 @@ index 296aa2b3..e011c8d5 100644 constexpr bool oob_check = true; for (; kb0 < kb0_stop-1; ++kb0) { constexpr bool last_iter = false; -@@ -1793,7 +1859,7 @@ static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_u +@@ -1794,7 +1860,7 @@ static constexpr __host__ __device__ bool ggml_cuda_flash_attn_ext_mma_f16_may_u (DKQ == 256 && DV == 256 && ncols1 == 8 && ncols2 == 8); } @@ -221,7 +225,7 @@ index 296aa2b3..e011c8d5 100644 __launch_bounds__(ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_mma_get_occupancy(DKQ, DV, ncols1*ncols2)) static __global__ void flash_attn_ext_f16( const char * Q_ptr, -@@ -1936,12 +2002,12 @@ static __global__ void flash_attn_ext_f16( +@@ -1937,12 +2003,12 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. if (kb0_start == 0) { constexpr bool needs_fixup = false; // CUDA block is working on an entire tile. @@ -236,7 +240,7 @@ index 296aa2b3..e011c8d5 100644 (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); } -@@ -1985,7 +2051,7 @@ static __global__ void flash_attn_ext_f16( +@@ -1986,7 +2052,7 @@ static __global__ void flash_attn_ext_f16( constexpr bool is_fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. constexpr bool needs_fixup = false; @@ -245,10 +249,23 @@ index 296aa2b3..e011c8d5 100644 (Q_f2, K_h2, V_h2, mask_h, indices, sinks_f, dstk, dst_meta, scale, slope, logit_softcap, ne01, ne02, gqa_ratio, ne11, stride_Q1, stride_Q2, stride_K, stride_V, stride_mask, jt, zt_gqa, kb0_start, kb0_stop); #else -@@ -2004,6 +2070,47 @@ static __global__ void flash_attn_ext_f16( - +@@ -2006,10 +2072,10 @@ static __global__ void flash_attn_ext_f16( bool ggml_cuda_flash_attn_ext_mma_f16_shall_use_sparse(const int cc, const ggml_tensor * dst, const int ncols1, const int ncols2); + static size_t ggml_cuda_fattn_mma_shared_size( +- int DKQ, int DV, int ncols1, int ncols2, int cc, int warp_size) { ++ int DKQ, int DV, int ncols1, int ncols2, int cc, int warp_size, int nstages_override = -1) { + const int ncols = ncols1 * ncols2; + const auto config = ggml_cuda_fattn_mma_get_config(DKQ, DV, ncols, cc); +- const int nstages = ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc); ++ const int nstages = nstages_override >= 0 ? nstages_override : ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc); + const bool swz = ggml_cuda_fattn_mma_get_swizzled(DKQ, DV, ncols, cc); + const int stride_K = config.nbatch_K2 + (swz ? 0 : 4); + const int stride_V = DKQ == 576 ? stride_K : config.nbatch_V2 + (swz ? 0 : 4); +@@ -2021,6 +2087,47 @@ static size_t ggml_cuda_fattn_mma_shared_size( + return std::max(combine, config.Q_in_reg ? std::max(q, kv + mask) : q + kv + mask); + } + +// Whether the long-sequence variant (see ggml_cuda_fattn_mma_get_nstages_longseq) can and should be used for dst. +// GGML_CUDA_FA_LONGSEQ=0 disables it, =1 enables it on every NVIDIA GPU with cp.async, unset = per-architecture default. +static bool ggml_cuda_fattn_mma_longseq_arch_default(const int cc) { @@ -293,10 +310,11 @@ index 296aa2b3..e011c8d5 100644 template void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * KQV = dst; -@@ -2012,13 +2119,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml +@@ -2029,13 +2136,22 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml constexpr int ncols = ncols1 * ncols2; +- const int nthreads = ggml_cuda_fattn_mma_get_nthreads(DKQ, DV, ncols, cc); +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + constexpr bool longseq_instantiated = ncols2 == 1 && DKQ == DV && (DKQ == 64 || DKQ == 128); + const bool longseq = longseq_instantiated && ggml_cuda_fattn_mma_use_longseq(cc, dst); @@ -305,19 +323,20 @@ index 296aa2b3..e011c8d5 100644 + const bool longseq = false; +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + - const int nthreads = ggml_cuda_fattn_mma_get_nthreads (DKQ, DV, ncols, cc); - const int nbatch_fa = ggml_cuda_fattn_mma_get_nbatch_fa (DKQ, DV, ncols, cc); - const int nbatch_K2 = ggml_cuda_fattn_mma_get_nbatch_K2 (DKQ, DV, ncols, cc); - const int nbatch_V2 = ggml_cuda_fattn_mma_get_nbatch_V2 (DKQ, DV, ncols, cc); - const int nbatch_combine = ggml_cuda_fattn_mma_get_nbatch_combine(DKQ, DV, ncols, cc); - const bool Q_in_reg = ggml_cuda_fattn_mma_get_Q_in_reg (DKQ, DV, ncols, cc); -- const int nstages = ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc); -+ const int nstages = longseq ? ggml_cuda_fattn_mma_get_nstages_longseq(DKQ, DV, ncols, cc) : -+ ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc); - - const int cols_per_warp = std::min(ncols, get_cols_per_warp(cc)); ++ const int nthreads = ggml_cuda_fattn_mma_get_nthreads (DKQ, DV, ncols, cc); + const int nbatch_fa = ggml_cuda_fattn_mma_get_nbatch_fa(DKQ, DV, ncols, cc); +- const int nstages = ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc); ++ const int nstages = longseq ? ggml_cuda_fattn_mma_get_nstages_longseq(DKQ, DV, ncols, cc) : ++ ggml_cuda_fattn_mma_get_nstages (DKQ, DV, ncols1, ncols2, cc); const int warp_size_host = ggml_cuda_info().devices[ctx.device].warp_size; -@@ -2052,6 +2168,40 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml + const int nwarps = nthreads / warp_size_host; + constexpr bool V_is_K_view = DKQ == 576; // Guaranteed by the kernel selection logic in fattn.cu +- const size_t nbytes_shared_total = ggml_cuda_fattn_mma_shared_size(DKQ, DV, ncols1, ncols2, cc, warp_size_host); ++ const size_t nbytes_shared_total = ggml_cuda_fattn_mma_shared_size(DKQ, DV, ncols1, ncols2, cc, warp_size_host, nstages); + + float logit_softcap; + memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float)); +@@ -2047,6 +2163,37 @@ void ggml_cuda_flash_attn_ext_mma_f16_case(ggml_backend_cuda_context & ctx, ggml #endif // defined(GGML_USE_HIP) fattn_kernel_t fattn_kernel; bool use_sparse = false; @@ -334,10 +353,7 @@ index 296aa2b3..e011c8d5 100644 + if (ncols1 <= 64) { + // Same grid as the stock kernel of this tile size: identical stream-k seams, so results are bit-identical. + const int nstages_stock = ggml_cuda_fattn_mma_get_nstages(DKQ, DV, ncols1, ncols2, cc); -+ const size_t nbytes_KV_stock = nstages_stock <= 1 ? nbytes_shared_KV_1stage : nbytes_shared_KV_2stage; -+ const size_t nbytes_total_stock = std::max(nbytes_shared_combine, Q_in_reg ? -+ std::max(nbytes_shared_Q, nbytes_KV_stock + nbytes_shared_mask) : -+ nbytes_shared_Q + nbytes_KV_stock + nbytes_shared_mask); ++ const size_t nbytes_total_stock = ggml_cuda_fattn_mma_shared_size(DKQ, DV, ncols1, ncols2, cc, warp_size_host, nstages_stock); + fattn_kernel_t stock_kernel = flash_attn_ext_f16; + static bool stock_limit_raised[GGML_CUDA_MAX_DEVICES] = {false}; + if (!stock_limit_raised[id]) { @@ -358,24 +374,23 @@ index 296aa2b3..e011c8d5 100644 if (logit_softcap == 0.0f) { constexpr bool use_logit_softcap = false; #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) -@@ -2145,6 +2295,12 @@ DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(112, 112, 64) - DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 64) +@@ -2141,8 +2288,10 @@ DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(128, 128, 64) DECL_FATTN_MMA_F16_CASE_ALL_NCOLS2(256, 256, 64) -+#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) -+// Long-sequence variant only: + #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) ++// Wide tile (fattn-mma-wide.cu) and long-sequence variant: +extern DECL_FATTN_MMA_F16_CASE( 64, 64, 128, 1); -+extern DECL_FATTN_MMA_F16_CASE(128, 128, 128, 1); + extern DECL_FATTN_MMA_F16_CASE(128, 128, 128, 1); +-#endif +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) -+ + extern DECL_FATTN_MMA_F16_CASE(512, 512, 4, 2); extern DECL_FATTN_MMA_F16_CASE(512, 512, 8, 2); - extern DECL_FATTN_MMA_F16_CASE(512, 512, 16, 2); diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu -index 5fa3e530..9fb380e8 100644 +index 99e25cbd..b22d7bcd 100644 --- a/src/ggml-cuda/fattn.cu +++ b/src/ggml-cuda/fattn.cu -@@ -187,6 +187,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con +@@ -197,6 +197,15 @@ static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_con return; } @@ -393,23 +408,23 @@ index 5fa3e530..9fb380e8 100644 diff --git a/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu b/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu new file mode 100644 -index 00000000..b9bf54a1 +index 00000000..a255a4b0 --- /dev/null +++ b/src/ggml-cuda/template-instances/fattn-mma-f16-instance-ncols1_128-ncols2_1.cu @@ -0,0 +1,8 @@ +// Long-sequence FlashAttention variant (CUDA only), see ggml_cuda_fattn_mma_use_longseq in fattn-mma-f16.cuh. ++// The DKQ == 128 ncols1 == 128 case is instantiated by fattn-mma-wide.cu. + +#include "../fattn-mma-f16.cuh" + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +DECL_FATTN_MMA_F16_CASE(64, 64, 128, 1); -+DECL_FATTN_MMA_F16_CASE(128, 128, 128, 1); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index 5e2db805..83136355 100644 +index 0f990e05..3ab8ad0a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -11710,6 +11710,21 @@ static std::vector> make_test_cases_eval() { +@@ -11721,6 +11721,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(65, 67, 4, {1, 1}, 17, 75, false, false, 0, 1.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(65, 67, 4, {1, 1}, 113, 75, false, false, 0, 1.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); diff --git a/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch index 0d086f6214..408d5a2ffc 100644 --- a/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch +++ b/scripts/unsloth/ggml-patches/0003-ggml-cuda-h3-vae-epilogue-fusion.patch @@ -23,14 +23,14 @@ the 256-thread launch. src/ggml-cuda/norm.cu | 63 ++++++++++++ 2 files changed, 262 insertions(+), 1 deletion(-) -[Ported onto leejet/ggml dc998771 (stable-diffusion.cpp master-982): the diff below is +[Ported onto leejet/ggml 135ce1e0 (stable-diffusion.cpp master-1050): the diff below is the whole patch as one diff, regenerated from the port.] --- diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 8a295c19..74f3373a 100644 +index c926338b..eca69cae 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -2061,8 +2061,73 @@ struct batched_mul_mat_traits { +@@ -1737,8 +1737,73 @@ struct batched_mul_mat_traits { static inline auto convert_nc(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } }; @@ -105,7 +105,7 @@ index 8a295c19..74f3373a 100644 using traits = batched_mul_mat_traits; using cuda_t = typename traits::cuda_type; -@@ -2270,8 +2335,17 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const +@@ -1946,8 +2011,17 @@ static void ggml_cuda_mul_mat_cublas_impl(ggml_backend_cuda_context & ctx, const // Convert output back to F32 if needed if (cu_data_type != CUDA_R_32F) { @@ -123,8 +123,8 @@ index 8a295c19..74f3373a 100644 } } -@@ -4336,6 +4410,123 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, - } +@@ -4028,6 +4102,123 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, + // try and fuse nodes and return the number of nodes to skip +static bool ggml_cuda_cublas_epilogue_enabled() { @@ -247,7 +247,7 @@ index 8a295c19..74f3373a 100644 static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { if (ggml_cuda_fusion_disabled()) { -@@ -5083,6 +5274,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4776,6 +4967,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 2; } diff --git a/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch index 386d09a04b..39eae2c75b 100644 --- a/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch +++ b/scripts/unsloth/ggml-patches/0004-ggml-h3-dit-fused-ops.patch @@ -42,26 +42,30 @@ partial and full rotation, F32/F16 outputs, chained segments and the scales. create mode 100644 src/ggml-cuda/rope-pe.cu create mode 100644 src/ggml-cuda/rope-pe.cuh -[Ported onto leejet/ggml dc998771 (stable-diffusion.cpp master-982): the diff below is - the whole patch as one diff, regenerated from the port.] +[Ported onto leejet/ggml 135ce1e0 (stable-diffusion.cpp master-1050): the diff below is + the whole patch as one diff, regenerated from the port. + Upstream ggml 394d41dd, 62da2f1b, 50ebf58e, 9201d070 moved ROPE_APPLY, SAGE_ATTN, SOL_ATTN + and QUANTIZE_I8_CONVROT out of ggml; the ops added here (ROPE_PE_PERMUTE, MODULATE_ROWS, + SWIGLU_SCALED) now follow GLU directly and GGML_OP_COUNT is 104.] --- diff --git a/include/ggml.h b/include/ggml.h -index 49a6751d..bbafe175 100644 +index 4efff011..85484f37 100644 --- a/include/ggml.h +++ b/include/ggml.h -@@ -607,6 +607,9 @@ extern "C" { - GGML_OP_SAGE_ATTN, - GGML_OP_SOL_ATTN, - GGML_OP_ROPE_APPLY, +@@ -603,6 +603,10 @@ extern "C" { + + GGML_OP_GLU, + + GGML_OP_ROPE_PE_PERMUTE, + GGML_OP_MODULATE_ROWS, + GGML_OP_SWIGLU_SCALED, - ++ GGML_OP_COUNT, }; -@@ -2556,6 +2559,65 @@ extern "C" { - float scale, - float tau); + +@@ -2502,6 +2506,65 @@ extern "C" { + float max_bias, + float logit_softcap); + // Rotary embedding with a precomputed cos/sin table, fused with the head-major relayout + // attention kernels consume, an optional scale and an optional cast to F16. @@ -126,24 +130,24 @@ index 49a6751d..bbafe175 100644 struct ggml_tensor * a, enum ggml_prec prec), diff --git a/src/ggml-backend-meta.cpp b/src/ggml-backend-meta.cpp -index fb447e97..899c8ad4 100644 +index 679d5216..d2a1904f 100644 --- a/src/ggml-backend-meta.cpp +++ b/src/ggml-backend-meta.cpp -@@ -153,7 +153,8 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( +@@ -153,6 +153,9 @@ static ggml_backend_buffer_type_t ggml_backend_meta_device_get_host_buffer_type( static bool ggml_backend_meta_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { GGML_ASSERT(ggml_backend_dev_is_meta(dev)); -- if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_SOL_ATTN) { -+ if (op->op == GGML_OP_SAGE_ATTN || op->op == GGML_OP_SOL_ATTN || op->op == GGML_OP_ROPE_PE_PERMUTE || -+ op->op == GGML_OP_MODULATE_ROWS || op->op == GGML_OP_SWIGLU_SCALED) { - return false; - } ++ if (op->op == GGML_OP_ROPE_PE_PERMUTE || op->op == GGML_OP_MODULATE_ROWS || op->op == GGML_OP_SWIGLU_SCALED) { ++ return false; ++ } const ggml_backend_meta_device_context * meta_dev_ctx = (const ggml_backend_meta_device_context *) dev->context; + return std::all_of(meta_dev_ctx->simple_devs.begin(), meta_dev_ctx->simple_devs.end(), + [op](ggml_backend_dev_t simple_dev) { return ggml_backend_dev_supports_op(simple_dev, op); }); diff --git a/src/ggml-cpu/ggml-cpu.c b/src/ggml-cpu/ggml-cpu.c -index fd55589f..30b8cbb5 100644 +index a5c40876..6e606b41 100644 --- a/src/ggml-cpu/ggml-cpu.c +++ b/src/ggml-cpu/ggml-cpu.c -@@ -2220,6 +2220,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm +@@ -2063,6 +2063,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_glu(params, tensor); } break; @@ -162,9 +166,9 @@ index fd55589f..30b8cbb5 100644 case GGML_OP_GET_REL_POS: { ggml_compute_forward_get_rel_pos(params, tensor); -@@ -2554,6 +2566,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { +@@ -2391,6 +2403,9 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + case GGML_OP_ROPE: case GGML_OP_ROPE_BACK: - case GGML_OP_ROPE_APPLY: case GGML_OP_ADD_REL_POS: + case GGML_OP_ROPE_PE_PERMUTE: + case GGML_OP_MODULATE_ROWS: @@ -173,7 +177,7 @@ index fd55589f..30b8cbb5 100644 n_tasks = n_threads; } break; diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp -index e8ef14f2..94333d30 100644 +index f49f72d7..64bafa97 100644 --- a/src/ggml-cpu/ops.cpp +++ b/src/ggml-cpu/ops.cpp @@ -12,6 +12,7 @@ @@ -184,8 +188,8 @@ index e8ef14f2..94333d30 100644 struct ggml_fp8_e4m3_cpu { uint8_t value; -@@ -12686,3 +12687,156 @@ void ggml_compute_forward_quantize_i8_convrot(const ggml_compute_params * params - memset(dst_data + rows * dst->nb[1], 0, (size_t)(rows_padded - rows) * dst->nb[1]); +@@ -12585,3 +12586,156 @@ void ggml_compute_forward_lightning_indexer( + } } } + @@ -342,10 +346,10 @@ index e8ef14f2..94333d30 100644 + } +} diff --git a/src/ggml-cpu/ops.h b/src/ggml-cpu/ops.h -index 19d7075e..d731a958 100644 +index 2728b08b..2e5b534e 100644 --- a/src/ggml-cpu/ops.h +++ b/src/ggml-cpu/ops.h -@@ -98,6 +98,9 @@ void ggml_compute_forward_win_part(const struct ggml_compute_params * params, st +@@ -97,6 +97,9 @@ void ggml_compute_forward_win_part(const struct ggml_compute_params * params, st void ggml_compute_forward_win_unpart(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_unary(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_glu(const struct ggml_compute_params * params, struct ggml_tensor * dst); @@ -356,10 +360,10 @@ index 19d7075e..d731a958 100644 void ggml_compute_forward_add_rel_pos(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rwkv_wkv6(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index cd0cee80..5c799ccd 100644 +index eca69cae..595f7bc0 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -47,6 +47,8 @@ +@@ -46,6 +46,8 @@ #include "ggml-cuda/rope.cuh" #include "ggml-cuda/roll.cuh" #include "ggml-cuda/scale.cuh" @@ -368,9 +372,9 @@ index cd0cee80..5c799ccd 100644 #include "ggml-cuda/snake.cuh" #include "ggml-cuda/softcap.cuh" #include "ggml-cuda/softmax.cuh" -@@ -3253,6 +3255,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg - case GGML_OP_SOL_ATTN: - ggml_cuda_sol_attn(ctx, dst); +@@ -2909,6 +2911,15 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg + case GGML_OP_FLASH_ATTN_EXT: + ggml_cuda_flash_attn_ext(ctx, dst); break; + case GGML_OP_ROPE_PE_PERMUTE: + ggml_cuda_op_rope_pe_permute(ctx, dst); @@ -384,10 +388,10 @@ index cd0cee80..5c799ccd 100644 case GGML_OP_CROSS_ENTROPY_LOSS: ggml_cuda_cross_entropy_loss(ctx, dst); break; -@@ -6784,6 +6795,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g - return ggml_cuda_sage_attn_supported(dev_ctx->device, op); - case GGML_OP_SOL_ATTN: - return ggml_cuda_sol_attn_supported(dev_ctx->device, op); +@@ -6636,6 +6647,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g + op->type == GGML_TYPE_F32; + case GGML_OP_FLASH_ATTN_EXT: + return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); + case GGML_OP_ROPE_PE_PERMUTE: + return ggml_cuda_rope_pe_permute_supported(op); + case GGML_OP_MODULATE_ROWS: @@ -704,41 +708,41 @@ index 00000000..b68b6af4 + +bool ggml_cuda_rope_pe_permute_supported(const ggml_tensor * op); diff --git a/src/ggml.c b/src/ggml.c -index da9ccada..a9f0b47d 100644 +index 524db8d9..44752bfe 100644 --- a/src/ggml.c +++ b/src/ggml.c -@@ -1129,9 +1129,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { - "SAGE_ATTN", - "SOL_ATTN", - "ROPE_APPLY", +@@ -1125,9 +1125,12 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + + "GLU", + + "ROPE_PE_PERMUTE", + "MODULATE_ROWS", + "SWIGLU_SCALED", }; --static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); -+static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT != 108"); +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); ++static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", -@@ -1249,9 +1252,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { - "sage_attn(x,y,z)", - "sol_attn(x,y,z)", - "rope_apply(x,pe)", +@@ -1241,9 +1244,12 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + + "glu(x)", + + "rope_pe_permute(x,pe)", + "modulate_rows(x,a,b)", + "swiglu_scaled(x,y)", }; --static_assert(GGML_OP_COUNT == 105, "GGML_OP_COUNT != 105"); -+static_assert(GGML_OP_COUNT == 108, "GGML_OP_COUNT != 108"); +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); ++static_assert(GGML_OP_COUNT == 104, "GGML_OP_COUNT != 104"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); -@@ -5721,6 +5727,125 @@ struct ggml_tensor* ggml_sol_attn( - return result; +@@ -5582,6 +5588,125 @@ struct ggml_tensor * ggml_flash_attn_ext( } + +struct ggml_tensor * ggml_rope_pe_permute( + struct ggml_context * ctx, + struct ggml_tensor * a, @@ -862,7 +866,7 @@ index da9ccada..a9f0b47d 100644 struct ggml_tensor * a, enum ggml_prec prec) { diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index 83136355..fdda5dc7 100644 +index 3ab8ad0a..de6ec4a6 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2192,6 +2192,118 @@ struct test_glu : public test_case { @@ -984,7 +988,7 @@ index 83136355..fdda5dc7 100644 struct test_glu_split : public test_case { const ggml_glu_op op; const ggml_type type; -@@ -9759,6 +9871,33 @@ static std::vector> make_test_cases_eval() { +@@ -9770,6 +9882,33 @@ static std::vector> make_test_cases_eval() { } } diff --git a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch index 73fb6cd4cc..410c109f82 100644 --- a/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch +++ b/scripts/unsloth/ggml-patches/0006-ggml-cuda-cudnn-attention.patch @@ -49,11 +49,14 @@ lengths, GQA with batch, F32 K/V, permuted Q) and unmasked perf cases. create mode 100644 src/ggml-cuda/cudnn-attn.cu create mode 100644 src/ggml-cuda/cudnn-attn.cuh -[Ported onto leejet/ggml dc998771 (stable-diffusion.cpp master-982): the diff below is - the whole patch as one diff, regenerated from the port.] +[Ported onto leejet/ggml 135ce1e0 (stable-diffusion.cpp master-1050): the diff below is + the whole patch as one diff, regenerated from the port. + Upstream ggml 62da2f1b moved SageAttention out of ggml (GGML_OP_SAGE_ATTN is gone), so + the SAGE_ATTN half of this patch (sage-attn.cu hook, GGML_CUDA_CUDNN_SAGE) is dropped; + cuDNN now only serves FLASH_ATTN_EXT.] --- diff --git a/CMakeLists.txt b/CMakeLists.txt -index 35252a83..da110ce1 100644 +index 2ed8f62d..1fe28c4d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -209,6 +209,9 @@ set (GGML_CUDA_FA_QUANTS "q4_0-q4_0;q8_0-q8_0;f16-f16;bf16-bf16" CACHE STRING @@ -434,10 +437,10 @@ index 00000000..9749f2ab + void * q, void * k, void * v, void * o, void * workspace); diff --git a/src/ggml-cuda/cudnn-attn.cu b/src/ggml-cuda/cudnn-attn.cu new file mode 100644 -index 00000000..a5b9de14 +index 00000000..a7e8911d --- /dev/null +++ b/src/ggml-cuda/cudnn-attn.cu -@@ -0,0 +1,384 @@ +@@ -0,0 +1,376 @@ +#include "cudnn-attn.cuh" + +#if defined(GGML_CUDA_USE_CUDNN) && !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) @@ -488,28 +491,20 @@ index 00000000..a5b9de14 + +static bool cudnn_attn_make_plan(int device, const ggml_tensor * dst, cudnn_attn_plan & plan) { + static const int attn_mode = cudnn_attn_env_mode("GGML_CUDA_CUDNN_ATTN"); -+ static const int sage_mode = cudnn_attn_env_mode("GGML_CUDA_CUDNN_SAGE"); + static const bool want_bf16 = cudnn_attn_env_set("GGML_CUDA_CUDNN_ATTN_BF16", false); -+ if (attn_mode == 0 || (dst->op == GGML_OP_SAGE_ATTN && sage_mode == 0)) { -+ return false; -+ } -+ if (dst->op != GGML_OP_FLASH_ATTN_EXT && dst->op != GGML_OP_SAGE_ATTN) { -+ return false; -+ } -+ // An explicitly requested SageAttention variant keeps its kernel; AUTO takes the cuDNN path. -+ if (dst->op == GGML_OP_SAGE_ATTN && ggml_get_op_params_i32(dst, 1) != GGML_SAGE_ATTN_AUTO) { ++ if (attn_mode == 0 || dst->op != GGML_OP_FLASH_ATTN_EXT) { + return false; + } + const int cc = ggml_cuda_info().devices[device].cc; + if (!GGML_CUDA_CC_IS_NVIDIA(cc) || cc < GGML_CUDA_CC_AMPERE) { + return false; + } -+ // Unset switches keep the existing kernels except where cuDNN measured faster at H3 shapes: against the flash -+ // attention kernel on sm89 and sm100 (slower on sm86), against the Sage kernel on sm100 (slower on sm86 and -+ // sm89). Other architectures were not measured. =1 opts in on any sm80+ GPU. ++ // Unset switches keep the existing kernels except where cuDNN measured faster at H3 shapes than the flash ++ // attention kernel: sm89 and sm100 (slower on sm86). Other architectures were not measured. =1 opts in on ++ // any sm80+ GPU. + const bool sm100 = cc >= 1000 && cc < GGML_CUDA_CC_BLACKWELL; + const bool sm89 = cc >= GGML_CUDA_CC_ADA_LOVELACE && cc < GGML_CUDA_CC_HOPPER; -+ if (dst->op == GGML_OP_SAGE_ATTN ? sage_mode < 0 && !sm100 : attn_mode < 0 && !sm89 && !sm100) { ++ if (attn_mode < 0 && !sm89 && !sm100) { + return false; + } + @@ -824,25 +819,23 @@ index 00000000..a5b9de14 +#endif diff --git a/src/ggml-cuda/cudnn-attn.cuh b/src/ggml-cuda/cudnn-attn.cuh new file mode 100644 -index 00000000..067bea63 +index 00000000..c4e98b04 --- /dev/null +++ b/src/ggml-cuda/cudnn-attn.cuh -@@ -0,0 +1,26 @@ +@@ -0,0 +1,24 @@ +#pragma once + +#include "common.cuh" + -+// cuDNN scaled dot product attention for GGML_OP_FLASH_ATTN_EXT and GGML_OP_SAGE_ATTN on NVIDIA ++// cuDNN scaled dot product attention for GGML_OP_FLASH_ATTN_EXT on NVIDIA +// sm80+ when the backend is built with GGML_CUDA_CUDNN and libcudnn 9 loads at run time. +// Only unmasked, sink-free, ALiBi-free, softcap-free attention with equal Q/K/V head sizes and at +// least 2^20 attention scores is taken; everything else, and every shape cuDNN rejects, keeps the +// existing kernels. +// -+// By default FLASH_ATTN_EXT takes it on sm89 and sm100 and SAGE_ATTN on sm100, where it measured faster. ++// By default it is taken on sm89 and sm100, where it measured faster than the flash attention kernels. +// Environment: +// GGML_CUDA_CUDNN_ATTN=0|1 never use cuDNN for attention / use it on every sm80+ GPU -+// GGML_CUDA_CUDNN_SAGE=0|1 keep the SageAttention kernels for GGML_OP_SAGE_ATTN / use cuDNN on every -+// sm80+ GPU (AUTO mode; an explicitly requested Sage variant keeps its kernel) +// GGML_CUDA_CUDNN_ATTN_BF16=1 run cuDNN in BF16 instead of F16 (F16 K/V are narrowed) +// GGML_CUDA_CUDNN_LIB=path library to open instead of libcudnn.so.9 +// GGML_CUDA_CUDNN_ATTN_LOG=1 log every call that takes (or falls back from) the cuDNN path @@ -855,7 +848,7 @@ index 00000000..067bea63 +// so the caller runs its own kernel. +bool ggml_cuda_cudnn_attn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/src/ggml-cuda/fattn.cu b/src/ggml-cuda/fattn.cu -index 9fb380e8..84c4e653 100644 +index b22d7bcd..d58d98a4 100644 --- a/src/ggml-cuda/fattn.cu +++ b/src/ggml-cuda/fattn.cu @@ -4,6 +4,7 @@ @@ -866,7 +859,7 @@ index 9fb380e8..84c4e653 100644 #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) // one list per group of ncols1 queries: a column is selected if any query of the group can see it -@@ -760,11 +761,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d +@@ -771,11 +772,15 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d const ggml_cuda_flash_attn_ext_f16_extra_data f16_extra = ggml_cuda_flash_attn_ext_get_f16_extra_data(dst, need_f16_K, need_f16_V); @@ -883,77 +876,11 @@ index 9fb380e8..84c4e653 100644 switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); -diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index 5c799ccd..f19aba44 100644 ---- a/src/ggml-cuda/ggml-cuda.cu -+++ b/src/ggml-cuda/ggml-cuda.cu -@@ -961,7 +961,7 @@ static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_t - ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor) - : ggml_nbytes(tensor); - if (tensor->op == GGML_OP_SAGE_ATTN) { -- size = ggml_cuda_sage_attn_get_alloc_size(tensor); -+ size = ggml_cuda_sage_attn_get_alloc_size(buft_ctx->device, tensor); - } - if (tensor->op == GGML_OP_SOL_ATTN) { - size = ggml_cuda_sol_attn_get_alloc_size(tensor); -diff --git a/src/ggml-cuda/sage-attn.cu b/src/ggml-cuda/sage-attn.cu -index 2c103eb0..ceb29ed2 100644 ---- a/src/ggml-cuda/sage-attn.cu -+++ b/src/ggml-cuda/sage-attn.cu -@@ -1,4 +1,5 @@ - #include "sage-attn.cuh" -+#include "cudnn-attn.cuh" - - #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && __CUDACC_VER_MAJOR__ >= 12 - -@@ -332,13 +333,18 @@ bool ggml_cuda_sage_attn_supported(int device, const ggml_tensor* op) { - op->ne[2] == q->ne[1] && op->ne[3] == q->ne[3]; - } - --size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op) { -+size_t ggml_cuda_sage_attn_get_alloc_size(int device, const ggml_tensor* op) { - // Scratch belongs to the graph allocation, so graph capture and VRAM planning see its full lifetime. -- return SageWorkspace(op).size; -+ // The cuDNN path, when it can run, reuses the same region for its own scratch. -+ return std::max(SageWorkspace(op).size, ggml_cuda_cudnn_attn_alloc_size(device, op)); - } - - void ggml_cuda_sage_attn(ggml_backend_cuda_context& ctx, ggml_tensor* dst) { - GGML_ASSERT(ggml_cuda_sage_attn_supported(ctx.device, dst)); -+ ggml_cuda_set_device(ctx.device); -+ if (ggml_cuda_cudnn_attn(ctx, dst)) { -+ return; -+ } - if (dst->src[0]->ne[0] == 64) { - sage_attn<64>(ctx, dst); - } else { -@@ -352,7 +358,7 @@ bool ggml_cuda_sage_attn_supported(int, const ggml_tensor*) { - return false; - } - --size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op) { -+size_t ggml_cuda_sage_attn_get_alloc_size(int, const ggml_tensor* op) { - return ggml_nbytes(op); - } - -diff --git a/src/ggml-cuda/sage-attn.cuh b/src/ggml-cuda/sage-attn.cuh -index 126c7a74..e2648857 100644 ---- a/src/ggml-cuda/sage-attn.cuh -+++ b/src/ggml-cuda/sage-attn.cuh -@@ -4,7 +4,7 @@ - #include "common.cuh" - - bool ggml_cuda_sage_attn_supported(int device, const ggml_tensor* op); --size_t ggml_cuda_sage_attn_get_alloc_size(const ggml_tensor* op); -+size_t ggml_cuda_sage_attn_get_alloc_size(int device, const ggml_tensor* op); - void ggml_cuda_sage_attn(ggml_backend_cuda_context& ctx, ggml_tensor* dst); - - #endif diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp -index 108da4a2..62be5695 100644 +index 70b348f4..a54d8745 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp -@@ -12008,6 +12008,21 @@ static std::vector> make_test_cases_eval() { +@@ -12019,6 +12019,21 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 10.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext_large_logits(256, 256, 4, {4, 1}, 1024, 75, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16)); @@ -975,7 +902,7 @@ index 108da4a2..62be5695 100644 test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3})); -@@ -12469,6 +12484,16 @@ static std::vector> make_test_cases_perf() { +@@ -12480,6 +12495,16 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {12, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 0)); } diff --git a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch index 79a61e5495..223b4b5720 100644 --- a/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch +++ b/scripts/unsloth/ggml-patches/0007-ggml-cuda-generic-dit-vae.patch @@ -1,15 +1,21 @@ +[Ported onto leejet/ggml 135ce1e0 (stable-diffusion.cpp master-1050): the diff below is + the whole patch as one diff, regenerated from the port. + Upstream ggml 394d41dd removed the builtin ROPE_APPLY op (stable-diffusion.cpp now runs + it as a backend extension, launched without PDL), so the rope_apply PDL race fix in + rope.cu is dropped.] +--- diff --git a/include/ggml.h b/include/ggml.h -index f7ef1573..bec0b7b6 100644 +index b35ca84f..d8b51c64 100644 --- a/include/ggml.h +++ b/include/ggml.h -@@ -2581,6 +2581,14 @@ extern "C" { +@@ -2528,6 +2528,14 @@ extern "C" { // ggml_rope_pe_permute of ggml_rms_norm(a, eps) (no weight): every head row of `a` is // RMS-normalised first, with the same arithmetic as GGML_OP_RMS_NORM on the same backend, so // the result matches the two-op chain bit for bit where the backend supports it. + // ggml_rope_pe_permute (or _rms) with interleaved pairing: rotated element i < n_rot pairs + // a[2*j] with a[2*j + 1] (j = i / 2, r = i % 2): y = a[2*j]*pe[0, r, j, t] + a[2*j + 1]*pe[1, r, j, t], -+ // the pairing ggml_rope_apply(interleaved = true) uses, with the same rounding (bit-identical -+ // to ggml_rope_apply -> ggml_scale -> ggml_cpy). Returns `t` (a ggml_rope_pe_permute result). ++ // the pairing of stable-diffusion.cpp's rope_apply extension op (interleaved = true), with the same ++ // rounding (bit-identical to rope_apply -> ggml_scale -> ggml_cpy). Returns `t` (a ggml_rope_pe_permute result). + GGML_API struct ggml_tensor * ggml_rope_pe_permute_set_interleaved( + struct ggml_tensor * t, + bool interleaved); @@ -18,10 +24,10 @@ index f7ef1573..bec0b7b6 100644 struct ggml_context * ctx, struct ggml_tensor * a, diff --git a/src/ggml-cpu/ops.cpp b/src/ggml-cpu/ops.cpp -index ebe6e202..f7d6b8f7 100644 +index 956ee5b1..908dfcd8 100644 --- a/src/ggml-cpu/ops.cpp +++ b/src/ggml-cpu/ops.cpp -@@ -12703,6 +12703,7 @@ void ggml_compute_forward_rope_pe_permute( +@@ -12602,6 +12602,7 @@ void ggml_compute_forward_rope_pe_permute( const float scale = ggml_get_op_params_f32(dst, 1); const bool norm = ggml_get_op_params_i32(dst, 2) != 0; const float eps = ggml_get_op_params_f32(dst, 3); @@ -29,7 +35,7 @@ index ebe6e202..f7d6b8f7 100644 const int half = n_rot / 2; const int64_t D = src0->ne[0]; -@@ -12741,17 +12742,28 @@ void ggml_compute_forward_rope_pe_permute( +@@ -12640,17 +12641,28 @@ void ggml_compute_forward_rope_pe_permute( if (n_rot > 0) { const float * t = (const float *) pe->data + l * (int64_t) half * 4; @@ -43,8 +49,8 @@ index ebe6e202..f7d6b8f7 100644 - const int j = i - r * half; - p1[i] = x[j + half] * t[4 * j + 2 * r + 1]; + if (inter) { -+ // the expression ggml_compute_forward_rope_apply (interleaved) evaluates, so the compiler -+ // contracts it (or not) the same way and the two ops stay bit-identical on this backend ++ // the expression the interleaved CPU rope_apply evaluates (stable-diffusion.cpp ++ // src/core/cpu/rope.cpp, formerly ggml_compute_forward_rope_apply), so the two stay bit-identical + for (int j = 0; j < half; ++j) { + const float a = x[2 * j]; + const float b = x[2 * j + 1]; @@ -69,10 +75,10 @@ index ebe6e202..f7d6b8f7 100644 if (D > n_rot) { memcpy(y.data() + n_rot, x + n_rot, (D - n_rot) * sizeof(float)); diff --git a/src/ggml-cuda/ggml-cuda.cu b/src/ggml-cuda/ggml-cuda.cu -index f19aba44..05ac1936 100644 +index 595f7bc0..23fe4a7f 100644 --- a/src/ggml-cuda/ggml-cuda.cu +++ b/src/ggml-cuda/ggml-cuda.cu -@@ -37,6 +37,9 @@ +@@ -36,6 +36,9 @@ #include "ggml-cuda/mmvq.cuh" #include "ggml-cuda/moe-weighted-reduction.cuh" #include "ggml-cuda/norm.cuh" @@ -82,7 +88,7 @@ index f19aba44..05ac1936 100644 #include "ggml-cuda/opt-step-adamw.cuh" #include "ggml-cuda/opt-step-sgd.cuh" #include "ggml-cuda/out-prod.cuh" -@@ -4538,6 +4541,106 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g +@@ -4230,6 +4233,106 @@ static int ggml_cuda_try_fuse_cublas_epilogue(ggml_backend_cuda_context & ctx, g return consumed; } @@ -189,7 +195,7 @@ index f19aba44..05ac1936 100644 static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { if (ggml_cuda_fusion_disabled()) { -@@ -4566,6 +4669,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4259,6 +4362,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } @@ -210,7 +216,7 @@ index f19aba44..05ac1936 100644 if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { -@@ -4578,6 +4695,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph +@@ -4271,6 +4388,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } @@ -1098,25 +1104,11 @@ index b6dd7a69..2fe9a599 100644 } } -diff --git a/src/ggml-cuda/rope.cu b/src/ggml-cuda/rope.cu -index 33cee75a..ef25d68f 100644 ---- a/src/ggml-cuda/rope.cu -+++ b/src/ggml-cuda/rope.cu -@@ -13,6 +13,9 @@ static __global__ void rope_apply_f32( - const float * x, const float * pe, float * dst, - int64_t half_dim, int64_t heads, int64_t tokens, int64_t pairs, - RopeApplyStrides strides) { -+ // Launched with programmatic dependent launch (ggml_cuda_kernel_launch): wait for the producer of x -+ // before reading it. Without this the kernel could read q/k while the preceding norm still wrote them. -+ ggml_cuda_pdl_sync(); - const int64_t index = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; - if (index >= pairs) { - return; diff --git a/src/ggml.c b/src/ggml.c -index 1dcc8a0f..ef3c750f 100644 +index 029f837d..c90bf325 100644 --- a/src/ggml.c +++ b/src/ggml.c -@@ -5773,6 +5773,14 @@ struct ggml_tensor * ggml_rope_pe_permute_rms( +@@ -5634,6 +5634,14 @@ struct ggml_tensor * ggml_rope_pe_permute_rms( return result; } diff --git a/scripts/unsloth/ggml-patches/0010-ggml-metal-rope-apply.patch b/scripts/unsloth/ggml-patches/0010-ggml-metal-rope-apply.patch deleted file mode 100644 index 280996ed28..0000000000 --- a/scripts/unsloth/ggml-patches/0010-ggml-metal-rope-apply.patch +++ /dev/null @@ -1,205 +0,0 @@ -diff --git a/src/ggml-metal/ggml-metal-device.cpp b/src/ggml-metal/ggml-metal-device.cpp -index 2c7972b0..1d2e0318 100644 ---- a/src/ggml-metal/ggml-metal-device.cpp -+++ b/src/ggml-metal/ggml-metal-device.cpp -@@ -2564,6 +2564,20 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_roll(ggml_metal_ - return res; - } - -+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rope_apply(ggml_metal_library_t lib, const ggml_tensor * op) { -+ assert(op->op == GGML_OP_ROPE_APPLY); -+ GGML_UNUSED(op); -+ -+ const char * name = "kernel_rope_apply_f32"; -+ -+ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); -+ if (!res.pipeline) { -+ res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); -+ } -+ -+ return res; -+} -+ - ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad(ggml_metal_library_t lib, const ggml_tensor * op) { - assert(op->op == GGML_OP_PAD); - -diff --git a/src/ggml-metal/ggml-metal-device.h b/src/ggml-metal/ggml-metal-device.h -index 4012b9d5..48d77221 100644 ---- a/src/ggml-metal/ggml-metal-device.h -+++ b/src/ggml-metal/ggml-metal-device.h -@@ -173,6 +173,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_upscale - struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad (ggml_metal_library_t lib, const struct ggml_tensor * op); - struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad_reflect_1d (ggml_metal_library_t lib, const struct ggml_tensor * op); - struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_roll (ggml_metal_library_t lib, const struct ggml_tensor * op); -+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rope_apply (ggml_metal_library_t lib, const struct ggml_tensor * op); - struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_arange (ggml_metal_library_t lib, const struct ggml_tensor * op); - struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_timestep_embedding(ggml_metal_library_t lib, const struct ggml_tensor * op); - struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_adamw (ggml_metal_library_t lib, const struct ggml_tensor * op); -diff --git a/src/ggml-metal/ggml-metal-device.m b/src/ggml-metal/ggml-metal-device.m -index 44519903..8634d284 100644 ---- a/src/ggml-metal/ggml-metal-device.m -+++ b/src/ggml-metal/ggml-metal-device.m -@@ -1695,6 +1695,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - return true; -+ case GGML_OP_ROPE_APPLY: -+ return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && -+ op->type == GGML_TYPE_F32 && ggml_is_contiguous(op) && -+ op->src[0]->ne[1] <= 65535 && op->src[0]->ne[3] <= 65535; - case GGML_OP_IM2COL: - return ggml_is_contiguous(op->src[1]) && op->src[1]->type == GGML_TYPE_F32 && (op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_F32); - case GGML_OP_CONV_2D: -diff --git a/src/ggml-metal/ggml-metal-impl.h b/src/ggml-metal/ggml-metal-impl.h -index 4ffb5a77..a358e659 100644 ---- a/src/ggml-metal/ggml-metal-impl.h -+++ b/src/ggml-metal/ggml-metal-impl.h -@@ -1177,6 +1177,22 @@ typedef struct { - int32_t s3; - } ggml_metal_kargs_roll; - -+typedef struct { -+ int32_t ne00; -+ int32_t ne01; -+ int32_t ne02; -+ int32_t ne03; -+ uint64_t nb00; -+ uint64_t nb01; -+ uint64_t nb02; -+ uint64_t nb03; -+ uint64_t nb10; -+ uint64_t nb11; -+ uint64_t nb12; -+ uint64_t nb13; -+ int32_t interleaved; -+} ggml_metal_kargs_rope_apply; -+ - typedef struct { - uint64_t nb1; - int dim; -diff --git a/src/ggml-metal/ggml-metal-ops.cpp b/src/ggml-metal/ggml-metal-ops.cpp -index 53c98161..551d62b0 100644 ---- a/src/ggml-metal/ggml-metal-ops.cpp -+++ b/src/ggml-metal/ggml-metal-ops.cpp -@@ -414,6 +414,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { - { - n_fuse = ggml_metal_op_rope(ctx, idx); - } break; -+ case GGML_OP_ROPE_APPLY: -+ { -+ n_fuse = ggml_metal_op_rope_apply(ctx, idx); -+ } break; - case GGML_OP_IM2COL: - { - n_fuse = ggml_metal_op_im2col(ctx, idx); -@@ -5257,6 +5261,47 @@ int ggml_metal_op_roll(ggml_metal_op_t ctx, int idx) { - return 1; - } - -+int ggml_metal_op_rope_apply(ggml_metal_op_t ctx, int idx) { -+ ggml_tensor * op = ctx->node(idx); -+ -+ ggml_metal_library_t lib = ctx->lib; -+ ggml_metal_encoder_t enc = ctx->enc; -+ -+ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); -+ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); -+ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); -+ -+ ggml_metal_kargs_rope_apply args = { -+ /*.ne00 =*/ ne00, -+ /*.ne01 =*/ ne01, -+ /*.ne02 =*/ ne02, -+ /*.ne03 =*/ ne03, -+ /*.nb00 =*/ nb00, -+ /*.nb01 =*/ nb01, -+ /*.nb02 =*/ nb02, -+ /*.nb03 =*/ nb03, -+ /*.nb10 =*/ nb10, -+ /*.nb11 =*/ nb11, -+ /*.nb12 =*/ nb12, -+ /*.nb13 =*/ nb13, -+ /*.interleaved =*/ ggml_get_op_params_i32(op, 0) != 0 ? 1 : 0, -+ }; -+ -+ auto pipeline = ggml_metal_library_get_pipeline_rope_apply(lib, op); -+ -+ const int nth = std::max(1, std::min(ne00/2, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline))); -+ -+ ggml_metal_encoder_set_pipeline(enc, pipeline); -+ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); -+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); -+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); -+ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); -+ -+ ggml_metal_encoder_dispatch_threadgroups(enc, ne02, ne01, ne03, nth, 1, 1); -+ -+ return 1; -+} -+ - int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { - ggml_tensor * op = ctx->node(idx); - -diff --git a/src/ggml-metal/ggml-metal-ops.h b/src/ggml-metal/ggml-metal-ops.h -index 583d1156..ef6be30f 100644 ---- a/src/ggml-metal/ggml-metal-ops.h -+++ b/src/ggml-metal/ggml-metal-ops.h -@@ -93,6 +93,7 @@ int ggml_metal_op_upscale (ggml_metal_op_t ctx, int idx); - int ggml_metal_op_pad (ggml_metal_op_t ctx, int idx); - int ggml_metal_op_pad_reflect_1d (ggml_metal_op_t ctx, int idx); - int ggml_metal_op_roll (ggml_metal_op_t ctx, int idx); -+int ggml_metal_op_rope_apply (ggml_metal_op_t ctx, int idx); - int ggml_metal_op_arange (ggml_metal_op_t ctx, int idx); - int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx); - int ggml_metal_op_argmax (ggml_metal_op_t ctx, int idx); -diff --git a/src/ggml-metal/kernels/rope.metal b/src/ggml-metal/kernels/rope.metal -index 401ceacb..b47563c6 100644 ---- a/src/ggml-metal/kernels/rope.metal -+++ b/src/ggml-metal/kernels/rope.metal -@@ -331,3 +331,46 @@ template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kerne - - template [[host_name("kernel_rope_vision_f32")]] kernel kernel_rope_vision_t kernel_rope_vision; - template [[host_name("kernel_rope_vision_f16")]] kernel kernel_rope_vision_t kernel_rope_vision; -+ -+// x [d, n_head, L, N] (any strides), pe [2, 2, d/2, L] of [[cos, -sin], [sin, cos]]; dst contiguous [d, L, n_head*N] -+kernel void kernel_rope_apply_f32( -+ constant ggml_metal_kargs_rope_apply & args, -+ device const char * src0, -+ device const char * src1, -+ device char * dst, -+ uint3 tgpig[[threadgroup_position_in_grid]], -+ uint3 tpitg[[thread_position_in_threadgroup]], -+ uint3 ntg[[threads_per_threadgroup]]) { -+ const int64_t token = tgpig.x; -+ const int64_t head = tgpig.y; -+ const int64_t batch = tgpig.z; -+ -+ const int64_t half_dim = args.ne00/2; -+ -+ device const char * x_row = src0 + head*args.nb01 + token*args.nb02 + batch*args.nb03; -+ device const char * pe_row = src1 + token*args.nb13; -+ device float * out = (device float *) dst + (token + args.ne02*(head + args.ne01*batch))*args.ne00; -+ -+ for (int64_t j = tpitg.x; j < half_dim; j += ntg.x) { -+ const int64_t i0 = args.interleaved ? 2*j : j; -+ const int64_t i1 = args.interleaved ? 2*j + 1 : j + half_dim; -+ -+ const float a = *(device const float *) (x_row + i0*args.nb00); -+ const float b = *(device const float *) (x_row + i1*args.nb00); -+ -+ device const char * m = pe_row + j*args.nb12; -+ -+ const float c = *(device const float *) (m); -+ const float s = *(device const float *) (m + args.nb11); -+ -+ if (args.interleaved) { -+ const float neg_s = *(device const float *) (m + args.nb10); -+ const float c1 = *(device const float *) (m + args.nb10 + args.nb11); -+ out[i0] = a*c + b*neg_s; -+ out[i1] = a*s + b*c1; -+ } else { -+ out[i0] = a*c - b*s; -+ out[i1] = a*s + b*c; -+ } -+ } -+} diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 7bbc383484..7c8b7492f9 100644 --- a/src/conditioning/conditioner.hpp +++ b/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" @@ -141,6 +142,7 @@ struct ConditionerParams { }; struct Conditioner { + virtual int resolve_clip_skip(int) const { return 0; } virtual ~Conditioner() = default; public: @@ -173,6 +175,10 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { std::vector token_embed_custom2; std::map> embedding_pos_map; + int resolve_clip_skip(int value) const override { + return value > 0 ? value : ((sd_version_is_sd2(version) || sd_version_is_sdxl(version)) ? 2 : 1); + } + FrozenCLIPEmbedderWithCustomWords(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, const std::map& orig_embedding_map, @@ -456,10 +462,6 @@ struct FrozenCLIPEmbedderWithCustomWords : public Conditioner { sd::Tensor hidden_states; // [n_token, hidden_size] or [n_token, hidden_size + hidden_size2] sd::Tensor pooled; - if (clip_skip <= 0) { - clip_skip = (sd_version_is_sd2(version) || sd_version_is_sdxl(version)) ? 2 : 1; - } - size_t chunk_len = 77; size_t chunk_count = tokens.size() / chunk_len; for (int chunk_idx = 0; chunk_idx < chunk_count; chunk_idx++) { @@ -651,6 +653,8 @@ struct SD3CLIPEmbedder : public Conditioner { std::shared_ptr clip_g; std::shared_ptr t5; + int resolve_clip_skip(int value) const override { return value > 0 ? value : 2; } + SD3CLIPEmbedder(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, std::shared_ptr weight_manager = nullptr, @@ -898,10 +902,6 @@ struct SD3CLIPEmbedder : public Conditioner { auto& t5_tokens = token_and_weights[2].first; auto& t5_weights = token_and_weights[2].second; - if (clip_skip <= 0) { - clip_skip = 2; - } - size_t chunk_len = 77; int64_t t0 = ggml_time_ms(); sd::Tensor hidden_states; @@ -1097,6 +1097,8 @@ struct FluxCLIPEmbedder : public Conditioner { std::shared_ptr t5; size_t chunk_len = 256; + int resolve_clip_skip(int value) const override { return value > 0 ? value : 2; } + FluxCLIPEmbedder(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, std::shared_ptr weight_manager = nullptr, @@ -1293,10 +1295,6 @@ struct FluxCLIPEmbedder : public Conditioner { auto& t5_tokens = token_and_weights[1].first; auto& t5_weights = token_and_weights[1].second; - if (clip_skip <= 0) { - clip_skip = 2; - } - int64_t t0 = ggml_time_ms(); sd::Tensor hidden_states; // [N, n_token, 4096] sd::Tensor pooled; // [768,] @@ -1833,11 +1831,11 @@ struct AnimaConditioner : public Conditioner { std::shared_ptr weight_manager = nullptr, const TokenizerConfig& tokenizers = {}) { llm = std::make_shared(LLM::LLMArch::QWEN3, - backend, - tensor_storage_map, - "text_encoders.llm", - false, - weight_manager); + backend, + tensor_storage_map, + "text_encoders.llm", + false, + weight_manager); qwen_tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 151643); if (!qwen_tokenizer) { qwen_tokenizer = std::make_shared(); @@ -1979,6 +1977,7 @@ struct LLMEmbedder : public Conditioner { std::shared_ptr tokenizer; std::shared_ptr llm; std::shared_ptr byt5; + std::shared_ptr ph3_connector; LLMEmbedder(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -2008,7 +2007,8 @@ struct LLMEmbedder : public Conditioner { arch = LLM::LLMArch::GPT_OSS_20B; } else if (sd_version_is_pid(version)) { arch = LLM::LLMArch::GEMMA2_2B; - } else if (version == VERSION_QWEN_IMAGE_2_1 || + } else if (version == VERSION_FIBO2 || + version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_flux3(version) || sd_version_is_lingbot_video(version) || sd_version_is_ideogram4(version) || @@ -2024,12 +2024,18 @@ struct LLMEmbedder : public Conditioner { } else if (version == VERSION_MOSS_IMAGE) { arch = LLM::LLMArch::QWEN3_5; } - llm = std::make_shared(arch, + llm = std::make_shared(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(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; @@ -2065,6 +2071,9 @@ struct LLMEmbedder : public Conditioner { void get_param_tensors(std::map& 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"); } @@ -2076,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); } @@ -2083,6 +2095,9 @@ struct LLMEmbedder : public Conditioner { void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); + if (ph3_connector) { + ph3_connector->set_runtime_backends(backends); + } if (byt5) { byt5->set_runtime_backends(backends); } @@ -2115,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); } @@ -2122,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& adapter) override { + if (ph3_connector) { + ph3_connector->set_weight_adapter(adapter); + } if (llm) { llm->set_weight_adapter(adapter); } @@ -2137,6 +2161,9 @@ struct LLMEmbedder : public Conditioner { } void runner_end() override { + if (ph3_connector) { + ph3_connector->runner_end(); + } if (llm) { llm->runner_end(); } @@ -2147,10 +2174,11 @@ struct LLMEmbedder : public Conditioner { std::tuple, std::vector, std::vector> tokenize(std::string text, const std::pair& attn_range, - size_t min_length = 0, - size_t max_length = 100000000, - bool spell_quotes = false, - const std::string& suffix = "") { + size_t min_length = 0, + size_t max_length = 100000000, + bool spell_quotes = false, + const std::string& suffix = "", + const std::vector>>& image_embeds = {}) { std::vector> parsed_attention; if (attn_range.first >= 0 && attn_range.second > 0) { if (attn_range.first > 0) { @@ -2195,6 +2223,19 @@ struct LLMEmbedder : public Conditioner { weights.insert(weights.end(), curr_tokens.size(), curr_weight); } + for (const auto& image_embed : image_embeds) { + const int index = image_embed.first; + const int64_t count = image_embed.second.shape()[1]; + if (index < 0 || static_cast(index) >= tokens.size() || count < 1) { + LOG_ERROR("Invalid reference image token range"); + return {}; + } + const int token = tokens[index]; + const float weight = weights[index]; + tokens.insert(tokens.begin() + index, count - 1, token); + weights.insert(weights.begin() + index, count - 1, weight); + } + if (!suffix.empty()) { std::vector suffix_tokens; if (!tokenizer->encode(suffix, suffix_tokens) || suffix_tokens.size() >= max_length) { @@ -2215,6 +2256,17 @@ struct LLMEmbedder : public Conditioner { std::vector mask; tokenizer->pad_tokens(tokens, &weights, &mask, min_length, max_length); + // Keep masks and weights in expanded coordinates, but only look up one token per image. + for (auto it = image_embeds.rbegin(); it != image_embeds.rend(); ++it) { + const int64_t begin = it->first; + const int64_t end = begin + it->second.shape()[1]; + if (begin < 0 || end > static_cast(tokens.size())) { + LOG_ERROR("Reference image exceeds the tokenized context"); + return {}; + } + tokens.erase(tokens.begin() + begin + 1, tokens.begin() + end); + } + // for (int i = 0; i < tokens.size(); i++) { // std::cout << tokens[i] << ":" << weights[i] << ", " << i << std::endl; // } @@ -2237,7 +2289,7 @@ struct LLMEmbedder : public Conditioner { const std::vector& image_grids = {}, const std::string& prompt_suffix = "", sd::Tensor* output_mask = nullptr) { - auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix); + auto tokens_weights_mask = tokenize(prompt, prompt_attn_range, min_length, max_length, spell_quotes, prompt_suffix, image_embeds); auto& tokens = std::get<0>(tokens_weights_mask); auto& weights = std::get<1>(tokens_weights_mask); auto& mask = std::get<2>(tokens_weights_mask); @@ -2267,10 +2319,10 @@ struct LLMEmbedder : public Conditioner { sd::Tensor position_ids; if (sd_version_is_flux3(version)) { - position_ids = sd::Tensor::zeros({static_cast(tokens.size() * 4)}); + position_ids = sd::Tensor::zeros({static_cast(weights.size() * 4)}); for (size_t axis = 0; axis < 4; ++axis) { - for (size_t i = 0; i < tokens.size(); ++i) { - position_ids[axis * tokens.size() + i] = mask[i] == 0.f ? 1 : static_cast(i); + for (size_t i = 0; i < weights.size(); ++i) { + position_ids[axis * weights.size() + i] = mask[i] == 0.f ? 1 : static_cast(i); } } } @@ -2288,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); @@ -2367,6 +2425,50 @@ struct LLMEmbedder : public Conditioner { } } + bool expand_vision_marker(const std::string& prompt, + size_t& marker_start, + std::vector> outputs, + int grid_h, + int grid_w, + std::vector>>& image_embeds, + LLM::DeepStackImageEmbeds& deepstack_image_embeds, + std::vector& image_grids, + int max_length = std::numeric_limits::max()) { + if (outputs.empty()) { + return false; + } + const std::string placeholder = "<|image_pad|>"; + size_t position = prompt.find(placeholder, marker_start); + if (position == std::string::npos) { + LOG_ERROR("Reference image marker is missing"); + return false; + } + std::vector prefix_tokens; + if (!tokenizer->encode(prompt.substr(0, position), prefix_tokens, nullptr)) { + return false; + } + int index = static_cast(prefix_tokens.size()); + for (const auto& image_embed : image_embeds) { + index += static_cast(image_embed.second.shape()[1]) - 1; + } + int count = static_cast(outputs[0].shape()[1]); + if (index > max_length || count > max_length - index) { + LOG_ERROR("Reference images exceed the %d-token context", max_length); + return false; + } + if (deepstack_image_embeds.empty()) { + deepstack_image_embeds.resize(outputs.size() - 1); + } + GGML_ASSERT(deepstack_image_embeds.size() == outputs.size() - 1); + image_embeds.emplace_back(index, std::move(outputs[0])); + for (size_t layer = 1; layer < outputs.size(); ++layer) { + deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer])); + } + image_grids.push_back({index, count, grid_h, grid_w}); + marker_start = position + placeholder.size(); + return true; + } + SDCondition get_learned_condition(int n_threads, const ConditionerParams& conditioner_params) override { std::string prompt; @@ -2390,13 +2492,13 @@ 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{13, 19, 25, 31, 37} : std::set{50}; + prompt_attn_range = {0, 0}; if (llm->enable_vision) { - const std::string placeholder = "<|image_pad|>"; - const int patch_size = llm->config.vision.patch_size; - const int factor = patch_size * llm->config.vision.spatial_merge_size; + const int patch_size = llm->config.vision.patch_size; + const int factor = patch_size * llm->config.vision.spatial_merge_size; auto resize_for_vision = [&](const sd::Tensor& image) { int height = static_cast(image.shape()[1]); @@ -2424,29 +2526,11 @@ struct LLMEmbedder : public Conditioner { int grid_h, int grid_w) { GGML_ASSERT(image_outputs.size() == 4); - auto image_embed = std::move(image_outputs[0]); prompt += "<|vision_start|>"; - std::vector prefix_tokens; - if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) { - return false; - } - int image_embed_idx = static_cast(prefix_tokens.size()); - image_embeds.emplace_back(image_embed_idx, image_embed); - if (deepstack_image_embeds.empty()) { - deepstack_image_embeds.resize(image_outputs.size() - 1); - } - for (size_t layer = 0; layer < deepstack_image_embeds.size(); ++layer) { - deepstack_image_embeds[layer].emplace_back(image_embed_idx, std::move(image_outputs[layer + 1])); - } - image_grids.push_back({image_embed_idx, - static_cast(image_embed.shape()[1]), - grid_h, - grid_w}); - for (int64_t i = 0; i < image_embed.shape()[1]; ++i) { - prompt += placeholder; - } - prompt += "<|vision_end|>"; - return true; + size_t marker_start = prompt.size(); + prompt += "<|image_pad|><|vision_end|>"; + return expand_vision_marker(prompt, marker_start, std::move(image_outputs), grid_h, grid_w, + image_embeds, deepstack_image_embeds, image_grids); }; const auto* references = conditioner_params.minimax_h3_references; @@ -2516,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 |subject_definitions:))"); + if (!std::regex_search(text, verbatim)) { + text = "Edit . " + text; + } else { + text = trim(text); + } + } + prompt += text; } else if (sd_version_is_hunyuan_video(version)) { prompt_template_encode_start_idx = 98; out_layers = {26}; @@ -2609,14 +2703,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } prompt += img_prompt; @@ -2660,26 +2753,14 @@ struct LLMEmbedder : public Conditioner { return {}; } prompt += (i == 0 ? "" : " ") + std::string("<|vision_start|>"; - std::vector prefix_tokens; - if (!tokenizer->encode(prompt, prefix_tokens, nullptr)) { + size_t marker_start = prompt.size(); + prompt += "<|image_pad|><|vision_end|>"; + if (!expand_vision_marker(prompt, marker_start, std::move(outputs), + static_cast(height) / llm->config.vision.patch_size, + static_cast(width) / llm->config.vision.patch_size, + image_embeds, deepstack_image_embeds, image_grids)) { return {}; } - int index = static_cast(prefix_tokens.size()); - int count = static_cast(outputs[0].shape()[1]); - image_embeds.emplace_back(index, std::move(outputs[0])); - if (deepstack_image_embeds.empty()) { - deepstack_image_embeds.resize(outputs.size() - 1); - } - for (size_t layer = 1; layer < outputs.size(); ++layer) { - deepstack_image_embeds[layer - 1].emplace_back(index, std::move(outputs[layer])); - } - image_grids.push_back({index, count, - static_cast(height) / llm->config.vision.patch_size, - static_cast(width) / llm->config.vision.patch_size}); - for (int j = 0; j < count; ++j) { - prompt += "<|image_pad|>"; - } - prompt += "<|vision_end|>"; } } prompt_attn_range.first = static_cast(prompt.size()); @@ -2690,7 +2771,6 @@ struct LLMEmbedder : public Conditioner { if (llm->enable_vision && conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("%s", sd_version_is_mage_flow(version) ? "MageFlowEditPipeline" : "QwenImageEditPlusPipeline"); prompt_template_encode_start_idx = 64; - int image_embed_idx = 64 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { @@ -2707,6 +2787,7 @@ struct LLMEmbedder : public Conditioner { } } + prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; std::string placeholder = "<|image_pad|>"; std::string img_prompt; @@ -2726,19 +2807,22 @@ struct LLMEmbedder : public Conditioner { auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); + std::string image_prefix = prompt + img_prompt + (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; + std::vector prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; img_prompt += (sd_version_is_mage_flow(version) ? "Image " : "Picture ") + std::to_string(i + 1) + ": <|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } - prompt = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; prompt += img_prompt; prompt_attn_range.first = static_cast(prompt.size()); @@ -2816,14 +2900,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } @@ -2888,14 +2971,13 @@ struct LLMEmbedder : public Conditioner { return {}; } int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); img_prompt += "Picture " + std::to_string(i + 1) + ": <|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(img_prompt.size() + static_cast(num_image_tokens) * placeholder.size() + 32); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } prompt += img_prompt; @@ -2913,7 +2995,6 @@ struct LLMEmbedder : public Conditioner { LOG_INFO("LongCatEditPipeline"); prompt_template_encode_start_idx = 67; min_length = 512 + prompt_template_encode_start_idx; - int image_embed_idx = 36 + 6; int min_pixels = conditioner_params.ref_image_params.vlm_min_size; if (min_pixels <= 0) { @@ -2930,6 +3011,7 @@ struct LLMEmbedder : public Conditioner { } } + prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n"; std::string placeholder = "<|image_pad|>"; std::string img_prompt; @@ -2948,19 +3030,22 @@ struct LLMEmbedder : public Conditioner { auto resized_image = clip_preprocess(image, w_bar, h_bar); auto image_embed = llm->encode_image(n_threads, resized_image, false); GGML_ASSERT(!image_embed.empty()); + std::string image_prefix = prompt + img_prompt + "<|vision_start|>"; + std::vector prefix_tokens; + if (!tokenizer->encode(image_prefix, prefix_tokens, nullptr)) { + return {}; + } + int image_embed_idx = static_cast(prefix_tokens.size()); + for (const auto& previous_image : image_embeds) { + image_embed_idx += static_cast(previous_image.second.shape()[1]) - 1; + } image_embeds.emplace_back(image_embed_idx, image_embed); - image_embed_idx += 1 + static_cast(image_embed.shape()[1]) + 6; img_prompt += "<|vision_start|>"; - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int j = 0; j < num_image_tokens; j++) { - img_prompt += placeholder; - } + img_prompt += placeholder; img_prompt += "<|vision_end|>"; } - prompt = "<|im_start|>system\nAs an image editing expert, first analyze the content and attributes of the input image(s). Then, based on the user's editing instructions, clearly and precisely determine how to modify the given image(s), ensuring that only the specified parts are altered and all other aspects remain consistent with the original(s).<|im_end|>\n<|im_start|>user\n"; prompt += img_prompt; } else { prompt_template_encode_start_idx = 36; @@ -3028,7 +3113,10 @@ struct LLMEmbedder : public Conditioner { } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { prompt_template_encode_start_idx = 0; out_layers = {35}; // -2 - + if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { + LOG_ERROR("Z-Image currently does not support reference-image conditioning"); + return {}; + } if (conditioner_params.ref_images != nullptr && !conditioner_params.ref_images->empty()) { LOG_INFO("ZImageOmniPipeline"); prompt = "<|im_start|>user\n<|vision_start|>"; @@ -3267,7 +3355,7 @@ struct LLMEmbedder : public Conditioner { for (const auto& [index, image_embed] : image_embeds) { int64_t begin = std::max(0, index - 1); int64_t end = std::min(static_cast(tags.size()), - index + image_embed.shape()[1] + 1); + index + image_embed.shape()[1] + 1); std::fill(tags.begin() + begin, tags.begin() + end, 0); } int64_t tag_count = static_cast(tags.size()); @@ -3447,9 +3535,9 @@ struct LLaDAImageEmbedder : public Conditioner { // <|endoftext|> doubles as the pad token in LLaDA2's tokenizer.json. tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 156892); query_former = std::make_shared(backend, - tensor_storage_map, - query_former_prefix, - weight_manager); + tensor_storage_map, + query_former_prefix, + weight_manager); text_projection = std::make_shared(backend, tensor_storage_map, text_projection_prefix, @@ -3563,12 +3651,10 @@ struct LLaDAImageEmbedder : public Conditioner { auto inputs_embeds = llm->compute_input_embeds(n_threads, text_ids); auto query_embeds = query_former->compute(n_threads, inputs_embeds); - // splice_image_embeds() replaces tokens in place, so the query slots have to exist in - // input_ids; their ids are irrelevant because the embeddings are overwritten. - std::vector padded(tokens.begin(), tokens.end()); - padded.resize(static_cast(n_text + num_queries), tokenizer->PAD_TOKEN_ID); - int64_t n_total = static_cast(padded.size()); - sd::Tensor input_ids({n_total}, padded); + std::vector input_tokens(tokens.begin(), tokens.end()); + input_tokens.push_back(tokenizer->PAD_TOKEN_ID); + sd::Tensor input_ids({n_text + 1}, std::move(input_tokens)); + int64_t n_total = n_text + num_queries; // Bidirectional everywhere except that the text tokens must not see the appended // queries, matching backbone_attention_mask[:, :, :text_length, text_length:] = min. @@ -3653,20 +3739,20 @@ struct LTXAVEmbedder : public Conditioner { LLM::LLMArch arch = detect_gemma_arch(tensor_storage_map, llm_prefix); LOG_INFO("ltxav text encoder: %s", arch == LLM::LLMArch::GEMMA4_12B ? "gemma 4" : "gemma 3"); llm = std::make_shared(arch, - backend, - tensor_storage_map, - llm_prefix, - false, - weight_manager); + backend, + tensor_storage_map, + llm_prefix, + false, + weight_manager); tokenizer = tokenizers.create(TokenizerConfig::MAIN, llm->config.vocab_size, 0, true); if (!tokenizer) { tokenizer = std::make_shared(); } dual_projection = tensor_storage_map.find(projector_prefix + ".video_aggregate_embed.weight") != tensor_storage_map.end(); projector = std::make_shared(backend, - tensor_storage_map, - projector_prefix, - weight_manager); + tensor_storage_map, + projector_prefix, + weight_manager); } void get_param_tensors(std::map& tensors) override { diff --git a/src/conditioning/conditioning_cache.h b/src/conditioning/conditioning_cache.h index 6e6cbcaf77..34648c9ad1 100644 --- a/src/conditioning/conditioning_cache.h +++ b/src/conditioning/conditioning_cache.h @@ -48,7 +48,8 @@ class ConditioningCache { return std::tie(p.text, p.clip_skip, p.width, p.height, p.zero_out_masked, r.pass_to_vlm, r.pass_to_dit, r.ref_index_mode, r.force_ref_timestep_zero, r.resize_before_vae, r.vae_input_max_pixels, - r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target); + r.vlm_resize_mode, r.vlm_min_size, r.vlm_max_size, r.resize_vae_to_target, + r.ref_image_size); }; if (fields(a) != fields(b) || (a.ref_images == nullptr) != (b.ref_images == nullptr) || diff --git a/src/conditioning/fibo2.hpp b/src/conditioning/fibo2.hpp new file mode 100644 index 0000000000..43b2ee0845 --- /dev/null +++ b/src/conditioning/fibo2.hpp @@ -0,0 +1,182 @@ +#ifndef __SD_CONDITIONING_FIBO2_HPP__ +#define __SD_CONDITIONING_FIBO2_HPP__ + +#include "conditioning/conditioner.hpp" + +struct Fibo2Conditioner : public LLMEmbedder { + using JSON = nlohmann::ordered_json; + + Fibo2Conditioner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr weight_manager, const TokenizerConfig& tokenizers) + : LLMEmbedder(backend, weights, VERSION_FIBO2, "", true, weight_manager, tokenizers) { + if (llm->config.hidden_size != 2560 || llm->config.num_layers != 36) { + throw std::runtime_error("Fibo 2.0 requires its Qwen3-VL-4B text encoder"); + } + } + + static bool empty_value(const JSON& value) { + return value.is_null() || + (value.is_string() && value.get_ref().empty()) || + ((value.is_array() || value.is_object()) && value.empty()); + } + + static JSON clean_caption(const JSON& value) { + if (value.is_object()) { + auto result = JSON::object(); + for (auto it = value.begin(); it != value.end(); ++it) { + auto cleaned = clean_caption(it.value()); + if (!empty_value(cleaned)) { + result[it.key()] = std::move(cleaned); + } + } + return result; + } + if (value.is_array()) { + auto result = JSON::array(); + for (const auto& item : value) { + auto cleaned = clean_caption(item); + if (!empty_value(cleaned)) { + result.push_back(std::move(cleaned)); + } + } + return result; + } + return value; + } + + static std::pair parse_prompt(const std::string& prompt) { + auto caption = JSON::parse(prompt, nullptr, false); + if (!caption.is_object()) { + return {prompt, nullptr}; + } + JSON instruction; + if (caption.contains("edit_instruction")) { + instruction = caption["edit_instruction"]; + caption.erase("edit_instruction"); + } + if (caption.contains("structured_caption")) { + caption = JSON(caption["structured_caption"]); + if (!caption.is_object()) { + return {caption, instruction}; + } + } + bool short_form = caption.size() == 2 && caption.contains("short_description") && caption.contains("text_render"); + caption = clean_caption(caption); + for (const char* score : {"aesthetic_score", "preference_score"}) { + caption.erase(score); + for (auto& value : caption) { + if (value.is_object()) { + value.erase(score); + } + } + } + caption = clean_caption(caption); + if (short_form) { + caption = JSON{{"short_description", caption.value("short_description", JSON(""))}, + {"text_render", caption.value("text_render", JSON::array())}}; + } + return {caption, instruction}; + } + + static std::string format_prompt(const std::string& text, int num_images) { + auto [caption, instruction] = parse_prompt(text); + if (num_images > 0) { + if (instruction.is_null()) { + caption = nullptr; + instruction = text; + } + auto edit = JSON::object(); + for (int i = 0; i < num_images; ++i) { + edit[num_images == 1 ? "image" : "image_" + std::to_string(i + 1)] = "<|vision_start|><|image_pad|><|vision_end|>"; + } + if (!empty_value(caption)) { + edit["structured_caption"] = caption; + } + edit["edit_instruction"] = instruction; + return edit.dump(); + } + if (text.empty()) { + return R"({"short_description":"","objects":[],"background_setting":"","lighting":{},"aesthetics":{},"photographic_characteristics":{},"style_medium":"","context":"","artistic_style":""})"; + } + if (caption.is_string()) { + return caption.get(); + } + return empty_value(caption) ? "{}" : caption.dump(); + } + + static sd::Tensor make_bundles(sd::Tensor states, const std::set& layers) { + constexpr int width = 2560; + constexpr int bundles[6][3] = {{9, 20, 31}, {5, 16, 27}, {7, 18, 29}, {9, 20, 31}, {11, 22, 33}, {13, 24, 35}}; + for (int64_t token = 0; token < states.shape()[1]; ++token) { + for (size_t layer = 0; layer < layers.size(); ++layer) { + float sum = 0.f; + int64_t offset = token * states.shape()[0] + layer * width; + for (int i = 0; i < width; ++i) { + sum += states[offset + i] * states[offset + i]; + } + float scale = 1.f / std::sqrt(sum / width + 1e-6f); + for (int i = 0; i < width; ++i) { + states[offset + i] *= scale; + } + } + } + sd::Tensor result({6 * 3 * width, states.shape()[1]}); + for (int b = 0; b < 6; ++b) { + for (int j = 0; j < 3; ++j) { + int index = static_cast(std::distance(layers.begin(), layers.find(bundles[b][j]))); + for (int64_t token = 0; token < states.shape()[1]; ++token) { + std::copy_n(states.data() + token * states.shape()[0] + index * width, width, + result.data() + token * result.shape()[0] + (b * 3 + j) * width); + } + } + } + return result; + } + + SDCondition get_learned_condition(int n_threads, const ConditionerParams& p) override { + int image_count = p.ref_images ? static_cast(p.ref_images->size()) : 0; + if (image_count && !llm->enable_vision) { + LOG_ERROR("Fibo 2.0 editing requires vision weights"); + return {}; + } + std::string prompt = "<|im_start|>user\n" + format_prompt(p.text, image_count) + "<|im_end|>\n<|im_start|>assistant\n"; + std::vector>> images; + LLM::DeepStackImageEmbeds deepstack; + std::vector grids; + size_t marker_start = 0; + for (int i = 0; i < image_count; ++i) { + const auto& image = (*p.ref_images)[i]; + int factor = llm->config.vision.patch_size * llm->config.vision.spatial_merge_size; + int height = static_cast(image.shape()[1]), width = static_cast(image.shape()[0]); + int h = std::max(factor, static_cast(std::round(static_cast(height) / factor)) * factor); + int w = std::max(factor, static_cast(std::round(static_cast(width) / factor)) * factor); + resize_image_dims(height, width, h, w, factor, + p.ref_image_params.vlm_min_size, + p.ref_image_params.vlm_max_size, + p.ref_image_params.vlm_resize_mode); + auto pixels = sd::ops::interpolate(image, {w, h, image.shape()[2], image.shape()[3]}, sd::ops::InterpolateMode::Bicubic, false, true); + auto rgb = sd::Tensor({w, h, 3, 1}); + for (int64_t pixel = 0; pixel < int64_t(w) * h; ++pixel) { + float alpha = pixels.shape()[2] == 4 ? pixels[pixel + 3 * w * h] : 1.f; + for (int channel = 0; channel < 3; ++channel) { + rgb[pixel + channel * w * h] = 2.f * std::clamp(pixels[pixel + channel * w * h] * alpha + 1.f - alpha, 0.f, 1.f) - 1.f; + } + } + auto outputs = llm->encode_image_outputs(n_threads, rgb, false); + if (!expand_vision_marker(prompt, marker_start, std::move(outputs), + h / llm->config.vision.patch_size, w / llm->config.vision.patch_size, + images, deepstack, grids, 4096)) { + return {}; + } + } + const std::set layers = {5, 7, 9, 11, 13, 16, 18, 20, 22, 24, 27, 29, 31, 33, 35}; + auto states = encode_prompt(n_threads, prompt, {0, 0}, 0, 0, images, layers, 0, false, 4096, deepstack, grids); + if (states.empty()) { + return {}; + } + SDCondition result; + result.c_crossattn = make_bundles(std::move(states), layers); + return result; + } +}; + +#endif // __SD_CONDITIONING_FIBO2_HPP__ diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index 38453def1c..1da926ee5f 100644 --- a/src/core/backend_fit.cpp +++ b/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; @@ -183,7 +185,7 @@ namespace sd::backend_fit { vm_statistics64_data_t stats{}; mach_msg_type_number_t count = HOST_VM_INFO64_COUNT; const bool ok = host_page_size(host, &page_size) == KERN_SUCCESS && - host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS; + host_statistics64(host, HOST_VM_INFO64, (host_info64_t)&stats, &count) == KERN_SUCCESS; mach_port_deallocate(mach_task_self(), host); if (ok) { return ((int64_t)stats.free_count + stats.inactive_count) * page_size; @@ -193,14 +195,23 @@ namespace sd::backend_fit { } static size_t select_main_device(const std::vector& 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& components, @@ -262,10 +273,10 @@ namespace sd::backend_fit { const auto& runtime_devices = plan.runtimes[ci].devices; const bool fits_runtime = !runtime_devices.empty() && - std::all_of(runtime_devices.begin(), runtime_devices.end(), [&](size_t di) { + std::all_of(runtime_devices.begin(), runtime_devices.end(), [&](size_t di) { const int64_t headroom = headroom_for(ci, di); return headroom <= remaining[di] && comp.params_bytes <= remaining[di] - headroom; - }); + }); if (fits_runtime) { decision.params_location = ParamsLocation::MAIN_GPU; decision.params_device = runtime_devices.front(); diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp index 53d80e2f05..f562416f79 100644 --- a/src/core/compute_workspace.cpp +++ b/src/core/compute_workspace.cpp @@ -8,11 +8,9 @@ #include #include +#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(ggml_backend_t backend) @@ -43,22 +41,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) { - ggml_tensor* node = ggml_graph_node(graph, i); - if (!ggml_backend_supports_op(backend_, node)) { - return true; - } - if (uses_blas() && node->op == GGML_OP_MUL_MAT && ggml_backend_supports_op(blas_backend_, node)) { - return true; - } - } - return false; - } - ggml_backend_sched_t ComputeWorkspace::make_scheduler(size_t graph_size) { std::vector backends{backend_}; backends.insert(backends.end(), extra_backends_.begin(), extra_backends_.end()); @@ -86,10 +68,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()) { @@ -97,7 +75,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; @@ -111,57 +89,31 @@ 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(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; - } - if (uses_blas()) { - static std::atomic logged{0}; - if (logged.fetch_add(1) < 4) { - int on_blas = 0; - for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { - on_blas += ggml_backend_sched_get_tensor_backend(scheduler_, ggml_graph_node(graph, i)) == blas_backend_; - } - LOG_INFO("cpu blas: %d of %d graph nodes run on %s", on_blas, ggml_graph_n_nodes(graph), - ggml_backend_name(blas_backend_)); - } - } - } 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; + const size_t capacity = static_cast(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; + } + if (uses_blas()) { + static std::atomic logged{0}; + if (logged.fetch_add(1) < 4) { + int on_blas = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + on_blas += ggml_backend_sched_get_tensor_backend(scheduler_, ggml_graph_node(graph, i)) == blas_backend_; } - reservation_ = std::move(signature); - ++reservations_; - } - if (!ggml_gallocr_alloc_graph(allocator_, graph)) { - release(); - return false; + LOG_INFO("cpu blas: %d of %d graph nodes run on %s", on_blas, ggml_graph_n_nodes(graph), + ggml_backend_name(blas_backend_)); } } active_ = true; @@ -170,12 +122,9 @@ namespace sd { ComputeWorkspace::Measurement ComputeWorkspace::measure( ggml_cgraph* graph, - size_t direct_bytes, const std::function& external_backend, - const AssignNodes& assign_nodes) { - if (!needs_scheduler(graph)) { - return {{{backend_, direct_bytes}}, false}; - } + const AssignNodes& assign_nodes, + bool cache_measurement) { std::vector tensors; std::unordered_set seen; auto visit = [&](const ggml_tensor* tensor) { @@ -183,11 +132,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); @@ -195,6 +144,51 @@ namespace sd { visit(source); } } + std::vector key; + if (cache_measurement) { + key.reserve(tensors.size() * sizeof(ggml_tensor)); + auto append = [&](const auto& value) { + const auto* bytes = reinterpret_cast(&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 { + const auto alignment_offset = reinterpret_cast(tensor->data) % ggml_backend_get_alignment(external); + append(alignment_offset); + } + 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) { @@ -223,22 +217,20 @@ namespace sd { ggml_backend_buffer_set_usage(buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); } copy->buffer = buffer; - copy->data = reinterpret_cast(static_cast(1)); - copy->extra = nullptr; + // Capability checks inspect alignment even before weights are resident. + copy->data = source->data ? source->data : reinterpret_cast(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]); - } + 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*>(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 sizes(extra_backends_.size() + 3); @@ -255,6 +247,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; } @@ -289,12 +285,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; } @@ -315,6 +310,6 @@ namespace sd { } return total; } - return allocator_ != nullptr && backend == backend_ ? ggml_gallocr_get_buffer_size(allocator_, 0) : 0; + return 0; } } diff --git a/src/core/compute_workspace.h b/src/core/compute_workspace.h index d1b8cf5bde..d74a16c3af 100644 --- a/src/core/compute_workspace.h +++ b/src/core/compute_workspace.h @@ -4,7 +4,6 @@ #include #include -#include "ggml-alloc.h" #include "ggml-backend.h" namespace sd { @@ -18,20 +17,17 @@ namespace sd { std::vector extra_backends_; ggml_backend_t cpu_backend_ = nullptr; ggml_backend_t blas_backend_ = nullptr; - ggml_gallocr_t allocator_ = nullptr; ggml_backend_sched_t scheduler_ = nullptr; size_t scheduler_capacity_ = 0; - std::vector reservation_; - bool active_ = false; - size_t reservations_ = 0; + bool active_ = false; + std::vector measurement_key_; + std::vector measured_buffers_; ggml_backend_sched_t make_scheduler(size_t graph_size); - bool needs_scheduler(ggml_cgraph* graph) const; public: struct Measurement { std::vector buffers; - bool scheduler = false; }; using AssignNodes = std::function; @@ -41,15 +37,14 @@ namespace sd { ComputeWorkspace& operator=(const ComputeWorkspace&) = delete; void set_extra_backends(const std::vector& 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& external_backend, - const AssignNodes& assign_nodes); + const AssignNodes& assign_nodes, + bool cache_measurement = false); void synchronize() const; bool segment_end() noexcept; bool release(); @@ -60,7 +55,6 @@ namespace sd { ggml_backend_t cpu_backend() const { return cpu_backend_; } void set_accelerator_threads(int n_threads) const; size_t bytes(ggml_backend_t backend) const; - size_t reservation_count() const { return reservations_; } }; } diff --git a/src/core/cpu/int8.cpp b/src/core/cpu/int8.cpp new file mode 100644 index 0000000000..ec4be835bd --- /dev/null +++ b/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 +#include +#include + +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(data); + const int64_t k = e.op.input->ne[0]; + const int group = e.op.group_size; + const auto* input = static_cast(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(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(data); + const auto& op = e.op; + const int64_t k = op.weight->ne[0], n = op.weight->ne[1]; + const auto* w = static_cast(op.weight->data); + const auto* ws = op.weight_scale ? static_cast(op.weight_scale->data) : nullptr; + const auto* bias = op.bias ? static_cast(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(dst->data), b = reinterpret_cast(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_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(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(scratch); + e.scales = reinterpret_cast(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(e.op.input->data); + e.scales = e.op.packed ? reinterpret_cast(e.quant + e.op.input->ne[0] * ((e.rows + 3) / 4 * 4)) : static_cast(e.op.input_scale->data); + } + bool alias = false; + for (auto* input : dst->src) + alias |= overlaps(dst, input); + e.output = alias ? reinterpret_cast(scratch + quant_size) : static_cast(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(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(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 diff --git a/src/core/cpu/modulate.cpp b/src/core/cpu/modulate.cpp new file mode 100644 index 0000000000..84a90ab68e --- /dev/null +++ b/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 + +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(static_cast(t->data) + offset); + } + + static void modulate_compute(ggml_tensor* dst, int ith, int nth, void*) { + auto* output = static_cast(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_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")); + GGML_ASSERT(run); + auto* dst = nodes[0]; + const auto address = reinterpret_cast(dst->data); + bool alias = false; + for (const auto* input : dst->src) { + if (input) { + const auto src = reinterpret_cast(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 diff --git a/src/core/cpu/registration.cpp b/src/core/cpu/registration.cpp new file mode 100644 index 0000000000..d63fd1f099 --- /dev/null +++ b/src/core/cpu/registration.cpp @@ -0,0 +1,13 @@ +#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_rope(device); + register_modulate(device); + } + +} +#endif diff --git a/src/core/cpu/registration.h b/src/core/cpu/registration.h new file mode 100644 index 0000000000..eba629c2a6 --- /dev/null +++ b/src/core/cpu/registration.h @@ -0,0 +1,15 @@ +#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); +} +#endif + +#endif diff --git a/src/core/cpu/rope.cpp b/src/core/cpu/rope.cpp new file mode 100644 index 0000000000..ad1ac79b24 --- /dev/null +++ b/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 +#include + +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(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(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3]; + const char* pos = static_cast(pe->data) + token * pe->nb[3]; + float* output = static_cast(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(input + col * x->nb[0]); + sum += static_cast(v * v); + } + const float mean = static_cast(sum / x->ne[0]); + scale = 1.0f / std::sqrt(mean + eps); + } + auto value = [&](int64_t col) { + float v = *reinterpret_cast(input + col * x->nb[0]); + if (execution.norm) { + const auto* weight = static_cast(execution.weight->data); + v = (v * scale) * *reinterpret_cast(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(m); + const float s = *reinterpret_cast(m + pe->nb[1]); + if (interleaved) { + const float neg_s = *reinterpret_cast(m + pe->nb[0]); + const float c1 = *reinterpret_cast(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(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_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 diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt new file mode 100644 index 0000000000..0a95d6cf7c --- /dev/null +++ b/src/core/cuda/CMakeLists.txt @@ -0,0 +1,85 @@ +cmake_policy(SET CMP0079 NEW) + +function(sd_cuda_sources VENDOR LANGUAGE) + set(SOURCES) + foreach(OP rope int8 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}/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") + 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() diff --git a/src/core/cuda/affine.cu b/src/core/cuda/affine.cu new file mode 100644 index 0000000000..2cd660b722 --- /dev/null +++ b/src/core/cuda/affine.cu @@ -0,0 +1,183 @@ +#include +#include +#include +#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(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(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 + 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(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(a->data), y = reinterpret_cast(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(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<<>>(static_cast(a->data), static_cast(b->data), + static_cast(residual->data), static_cast(output), a->ne[0] / 4, count4); + } else if (node->op == GGML_OP_CONCAT) { + const bool permuted = count == 3; + auto kernel = permuted ? concat_cast : concat_cast; + kernel<<>>(static_cast(a->data), static_cast(b->data), + static_cast(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1], + node->ne[permuted ? 1 : 2], count4); + } else { + scale_permute<<>>(static_cast(a->data), static_cast(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)); + } + +} diff --git a/src/core/cuda/bf16.cu b/src/core/cuda/bf16.cu new file mode 100644 index 0000000000..59f5073947 --- /dev/null +++ b/src/core/cuda/bf16.cu @@ -0,0 +1,277 @@ +#include "bf16.h" +#if defined(SD_CUDA_VENDOR_HIP) +#include +#elif defined(SD_CUDA_VENDOR_MUSA) +#include +#else +#include +#endif +#include +#include +#include +#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(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(a->data), y = reinterpret_cast(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(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(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(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(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(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_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(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(ggml_backend_buffer_get_base(execution->workspace)); + const Workspace work(nodes, count); + auto* packed = reinterpret_cast(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(gate) | reinterpret_cast(up)) % 16 == 0) { + pack_swiglu4<<>>(static_cast(gate), + static_cast(up), + reinterpret_cast(packed), size / 4, scale); + } else { + pack_swiglu<<>>(static_cast(gate), + static_cast(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)); + } +} diff --git a/src/core/cuda/bf16.h b/src/core/cuda/bf16.h new file mode 100644 index 0000000000..b5da3704f0 --- /dev/null +++ b/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 +#elif defined(SD_CUDA_VENDOR_MUSA) +#include +#else +#include +#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 diff --git a/src/core/cuda/bf16_gemm.cu b/src/core/cuda/bf16_gemm.cu new file mode 100644 index 0000000000..41d75706ca --- /dev/null +++ b/src/core/cuda/bf16_gemm.cu @@ -0,0 +1,25 @@ +#include +#include +#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(weight) | reinterpret_cast(input) | reinterpret_cast(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, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>; + Gemm gemm; + const Gemm::Arguments args({int(rows), int(n), int(k)}, + {static_cast(input), k}, + {static_cast(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 +} diff --git a/src/core/cuda/bf16_gemm.h b/src/core/cuda/bf16_gemm.h new file mode 100644 index 0000000000..6cb36501f9 --- /dev/null +++ b/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 +#include + +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 diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu new file mode 100644 index 0000000000..cf7896114d --- /dev/null +++ b/src/core/cuda/convrot.cu @@ -0,0 +1,306 @@ +#include +#include +#include +#include +#include +#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(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; + *workspace = ggml_nbytes(nodes[count - 1]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void* state, void*) { + delete static_cast(state); + } + static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + if (!b) + return false; + const auto x = reinterpret_cast(a->data), y = reinterpret_cast(b->data); + return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); + } + template + 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, + 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(t->data) : nullptr; }; + quantize<<>>( + data(op.input) + op.gate_offset, output, reinterpret_cast(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(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(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(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(plan, op, output, stream); + break; + case NORM: + error = launch(plan, op, output, stream); + break; + case RMS_MOD: + error = launch(plan, op, output, stream); + break; + case MOD: + error = launch(plan, op, output, stream); + break; + case SWIGLU: + case SILU: + error = launch(plan, op, output, stream); + break; + case SIGMOID: + error = launch(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)); + } + } + +} diff --git a/src/core/cuda/convrot_kernel.h b/src/core/cuda/convrot_kernel.h new file mode 100644 index 0000000000..04abf97ac3 --- /dev/null +++ b/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 +#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 + 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 + __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 diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu new file mode 100644 index 0000000000..cdafc9983e --- /dev/null +++ b/src/core/cuda/cute_attention.cu @@ -0,0 +1,792 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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 + __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 + __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 + __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 + __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{}); + return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l)); + } + }; + + template + __device__ __forceinline__ auto to_half(const TensorT& src) { + constexpr int count = decltype(size(src))::value; + auto dst = make_tensor(src.layout()); + cutlass::NumericArrayConverter convert; + *reinterpret_cast*>(dst.data()) = + convert(*reinterpret_cast*>(src.data())); + return dst; + } + + template + __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 + struct Softmax { + decltype(make_tensor(Shape>{})) maximum, sum; + template + __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 value = exp2f(scores(m, n) * 1.4426950408889634f - next_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 + __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 = 1.f / total; +#pragma unroll + for (int n = 0; n < size<1>(output); ++n) + output(m, n) *= inverse; + } + } + }; + + template + struct Traits { + using Element = cutlass::half_t; + static constexpr int Threads = Warps * 32; + using Mma = TiledMMA, + Layout, _1, _1>>, + Tile, _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{}, + Layout>, Stride, _1>>{})); + using QLayout = decltype(tile_to_shape(SmemAtom{}, Shape, Int>{})); + using KVLayout = decltype(tile_to_shape(SmemAtom{}, Shape, Int>{})); + using VtLayout = decltype(composition(KVLayout{}, make_layout(Shape, Int>{}, GenRowMajor{}))); + using CopyAtom = Copy_Atom; + using CopyAtomT = Copy_Atom; + using GlobalLayout = Layout, Int>, Stride, _1>>; + using GlobalCopy = decltype(make_tiled_copy(Copy_Atom, Element>{}, + GlobalLayout{}, + Layout>{})); + using OutputCopy = decltype(make_tiled_copy(Copy_Atom, Element>{}, + GlobalLayout{}, + Layout>{})); + 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; + }; + + template + __global__ __launch_bounds__(Warps * 32) void kernel(Params p) { + using T = Traits; + 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>{}, make_coord(m_block, 0)); + auto gK = local_tile(mK, Shape, Int>{}, make_coord(_, 0)); + auto gV = local_tile(mV, Shape, Int>{}, make_coord(_, 0)); + auto sQ = make_tensor(make_smem_ptr(reinterpret_cast(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>{})); + auto tKVcKV = thread_copy.partition_S(make_identity_tensor(Shape, Int>{})); + 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>{}); + 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>{})); + // Process the partial K/V tile first; subsequent tiles need no score masking. + int block = (p.nk - 1) / BN; + copy_rows(global_copy, tQgQ, tQsQ, tQcQ, p.nq - m_block * BM, p.dim); + auto last_k = tKgK(_, _, _, block); + copy_rows(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>{}); + clear(acc_s); + cute::cp_async_wait<0>(); + __syncthreads(); + auto v_tile = tVgV(_, _, _, n); + copy_rows(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(global_copy, next_k, tKsK, tKVcKV, 0, p.dim); + cute::cp_async_fence(); + } + softmax.template update(acc_s, acc_o); + auto rP = to_half(acc_s); + auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs(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, 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); + auto mO = make_tensor(make_gmem_ptr(p.o + int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim), + make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{})); + auto gO = local_tile(mO, Shape, Int>{}, make_coord(m_block, 0)); + typename T::OutputCopy output_copy; + auto global_threadO = output_copy.get_thread_slice(tid); + auto from_sO = global_threadO.partition_S(sO); + auto to_gO = global_threadO.partition_D(gO); + auto output_id = global_threadO.partition_D(make_identity_tensor(Shape, Int>{})); + __syncthreads(); + auto output_regs = make_tensor(shape(to_gO)); + cute::copy(output_copy, from_sO, output_regs); + copy_rows(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 +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(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(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 devices; + std::lock_guard 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_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(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(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(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(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(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]; +}; + +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]; + size_t size = 0; + for (int i = 0; i < 4; ++i) { + 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 +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 initialized; + { + std::lock_guard lock(mutex); + if (!initialized[context.device]) { + auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes); + if (error != cudaSuccess) + return error; + initialized[context.device] = true; + } + } + kernel<<(context.stream)>>>(p); + return cudaGetLastError(); +} + +template +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(context, p, batches) : attention_launch(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(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_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(context.stream); + auto* base = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + half* work[4]; + for (int i = 0; i < 4; ++i) + work[i] = reinterpret_cast(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(src->data); + } else { + auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half : copy_to_half; + kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast(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]; + sd_fattn_cute::Params p = {reinterpret_cast(inputs[0]), + reinterpret_cast(inputs[1]), reinterpret_cast(inputs[2]), + reinterpret_cast(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])}; + 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; + 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); + } + } + auto* dst = nodes[count - 1]; + 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(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<<>>(work[3], packed, + reinterpret_cast(packed + width * padded_rows), width, rows, scale, bias); + } else { + copy_from_half<<<(ggml_nelements(attn) + 255) / 256, 256, 0, stream>>>(work[3], static_cast(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(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)); +} diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu new file mode 100644 index 0000000000..21b6134904 --- /dev/null +++ b/src/core/cuda/i8_gated_residual.cu @@ -0,0 +1,461 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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 + struct TileConfig { + using Tile = cutlass::gemm::GemmShape; + using Warp = cutlass::gemm::GemmShape; + }; + + 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 + 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 + 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 + struct KernelTypes { + using Tile = typename Config::Tile; + using Warp = typename Config::Warp; + using Swizzle = std::conditional_t>; + using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Residual = cutlass::epilogue::threadblock::VisitorAuxLoad>; + using GateMul = cutlass::epilogue::threadblock::VisitorCompute; + using Add = cutlass::epilogue::threadblock::VisitorCompute; + using Scalar = cutlass::epilogue::threadblock::VisitorScalarBroadcast; + using ScaleOut = cutlass::epilogue::threadblock::Sm80EVT; + using Scaled = std::conditional_t; + using Gated = cutlass::epilogue::threadblock::Sm80EVT; + using WithResidual = cutlass::epilogue::threadblock::Sm80EVT; + using Output = WithResidual; + using Store = cutlass::epilogue::threadblock::VisitorAuxStore>; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + 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; + }; + + 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 cache; + std::lock_guard 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_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(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 + static typename Types::Gemm::Arguments arguments(const Operation& op, float* output, bool metadata) { + const auto* x = metadata ? nullptr : static_cast(op.mat.input->data); + const auto* w = metadata ? nullptr : static_cast(op.mat.weight->data); + const auto* xs = metadata ? nullptr : op.mat.packed ? reinterpret_cast(x + op.k * ((op.rows + 3) / 4 * 4)) + : static_cast(op.mat.input_scale->data); + const auto* ws = metadata ? nullptr : static_cast(op.mat.weight_scale->data); + const auto* gate = metadata ? nullptr : static_cast(op.gate->data); + auto* residual = metadata ? nullptr : static_cast(op.residual->data); + typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}}; + typename Types::Scaled::Arguments scaled; + if constexpr (std::is_same::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; + ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr; + }; + template + static ggml_status launch(Plan& plan, const Operation& op, float* output, void* scratch, cudaStream_t stream) { + using Kernel = typename Types::Kernel; + auto args = arguments(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, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess) + return GGML_STATUS_FAILED; + int occupancy = 0; + if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&occupancy, cutlass::Kernel2, + 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<<>>(params); + return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + template + static bool configure(const Operation& op, Plan& plan) { + using Types = KernelTypes; + using Kernel = typename Types::Kernel; + if (sizeof(typename Kernel::SharedStorage) > plan.device.shared) + return false; + auto args = arguments(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; + return true; + } + template + static bool configure(const Operation& op, Plan& plan) { + if (op.scale != 1.f) + return op.rows >= 8192 ? configure(op, plan) : configure(op, plan); + return op.rows >= 8192 ? configure(op, plan) : configure(op, plan); + } + static bool configure(const Operation& op, Plan& plan) { + auto try_tile = [&](Tile tile) { + switch (tile) { + case TILE_128_256: + return configure>(op, plan); + case TILE_128_128: + return configure>(op, plan); + case TILE_64_128: + return configure>(op, plan); + case TILE_16_128: + return configure>(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(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(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 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)); + if (!configure(describe(nodes, count), *plan)) + return GGML_STATUS_FAILED; + *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void* state, void*) { + delete static_cast(state); + } + static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + if (!b) + return false; + const auto x = reinterpret_cast(a->data), y = reinterpret_cast(b->data); + return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); + } + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) { + auto& plan = *static_cast(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_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(context.stream); + auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + bool alias = reinterpret_cast(dst->data) % 16 != 0; + 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 && input->data == dst->data) + continue; + alias |= overlaps(dst, input); + } + auto* output = alias ? scratch + plan.output_offset : static_cast(dst->data); + const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0; + auto status = plan.launch(plan, op, reinterpret_cast(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)); +} diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu new file mode 100644 index 0000000000..69f91d2aea --- /dev/null +++ b/src/core/cuda/i8_gemm.cu @@ -0,0 +1,330 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#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 + struct TileConfig { + using Tile = cutlass::gemm::GemmShape; + using Warp = cutlass::gemm::GemmShape; + }; + + 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 cache; + std::lock_guard 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_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(ld)), half_n(extent.strided() / 2) {} + CUTLASS_HOST_DEVICE WeightLayout(cutlass::MatrixCoord extent, int64_t ld) + : stride(static_cast(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 + struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore> { + using Base = cutlass::epilogue::threadblock::VisitorAuxStore>; + using Base::Base; + template + 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(&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(this->params_ptr->ptr_aux + offset) = value; + } + } + } + }; + template + CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) { + auto cb = Base::get_callbacks(tile, tid, shape); + return Callbacks(cb); + } + }; + + static __global__ void pair_scales(const float* src, float* dst, int n) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) + dst[i] = src[i / 2 + (i % 2) * (n / 2)]; + } + + template + struct DenseTypes { + using Tile = typename Config::Tile; + using Warp = typename Config::Warp; + using Swizzle = std::conditional_t>; + using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Store = cutlass::epilogue::threadblock::VisitorAuxStore>; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + 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; + }; + template + struct SwigluTypes { + using Tile = typename Config::Tile; + using Warp = typename Config::Warp; + using Swizzle = std::conditional_t>; + using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Store = SwigluStore; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT; + using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal, cutlass::epilogue::thread::LinearCombination, 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; + using Kernel = std::conditional_t, + cutlass::gemm::kernel::GemmWithEpilogueVisitor>; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + }; + struct Plan { + DeviceInfo device; + int64_t rows, n, k; + bool swiglu; + 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) = nullptr; + }; + template + 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, w, nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0}; + } + template + 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) { + using Kernel = typename Types::Kernel; + if ((reinterpret_cast(x) | reinterpret_cast(w) | reinterpret_cast(xs) | + reinterpret_cast(ws) | reinterpret_cast(out)) % + 16) + return GGML_STATUS_FAILED; + auto* paired = static_cast(scratch); + auto args = arguments(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, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess) + return GGML_STATUS_FAILED; + if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&p.occupancy, cutlass::Kernel2, 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 (params.get_workspace_size() > p.kernel_bytes) + return GGML_STATUS_FAILED; + if (p.swiglu) { + pair_scales<<>>(ws, paired, int(p.n)); + if (cudaGetLastError() != cudaSuccess) + return GGML_STATUS_FAILED; + } + if (params.init_workspace(static_cast(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess) + return GGML_STATUS_FAILED; + cutlass::arch::synclog_setup(); + cutlass::Kernel2<<>>(params); + return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + template + static bool configure_kernel(Plan& p) { + using Kernel = typename Types::Kernel; + if (sizeof(typename Kernel::SharedStorage) > p.device.shared) + return false; + auto args = arguments(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; + return true; + } + template + static bool configure_tile(Plan& p) { + if (p.swiglu) + return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); + return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); + } + static bool configure(Plan& p) { + auto try_tile = [&](Tile tile) { + switch (tile) { + case TILE_128_256: + return configure_tile>(p); + case TILE_128_128: + return configure_tile>(p); + case TILE_64_128: + return configure_tile>(p); + case TILE_16_128: + return configure_tile>(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 {}; +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) { +#ifdef SD_CUDA_FORCE_CUBLAS + return nullptr; +#else + const auto info = sd_i8_cutlass::device_info(device); + if (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; + if (!sd_i8_cutlass::configure(*p)) { + delete p; + return nullptr; + } + *workspace = p->scratch_offset + p->kernel_bytes; + return p; +#endif +} +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); +} +void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) { + delete plan; +} diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h new file mode 100644 index 0000000000..189df6ef0b --- /dev/null +++ b/src/core/cuda/i8_gemm.h @@ -0,0 +1,9 @@ +#ifndef __SD_CORE_CUDA_I8_GEMM_H__ +#define __SD_CORE_CUDA_I8_GEMM_H__ +#include +#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); +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); +void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan); +#endif diff --git a/src/core/cuda/i8_projection.cu b/src/core/cuda/i8_projection.cu new file mode 100644 index 0000000000..2f8ec3ff9e --- /dev/null +++ b/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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +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 +struct I8RopeStore : cutlass::epilogue::threadblock::VisitorAuxStore> { + using Base = cutlass::epilogue::threadblock::VisitorAuxStore>; + 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 + 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 + 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(&src(i)), vb = *reinterpret_cast(&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::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 + CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) { + auto cb = Base::get_callbacks(tile, tid, shape); + return Callbacks(cb, extra); + } +}; + +template +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; + using Col = cutlass::epilogue::threadblock::VisitorColBroadcast>; + using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast>; + using Mul = cutlass::epilogue::threadblock::VisitorCompute; + using ScaleX = cutlass::epilogue::threadblock::Sm80EVT; + using ScaleW = cutlass::epilogue::threadblock::Sm80EVT; + using Store = I8RopeStore; + using EVT = cutlass::epilogue::threadblock::Sm80EVT; + 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; +}; + +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 cache; + std::lock_guard 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_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::Kernel::SharedStorage), + sizeof(typename ProjectionKernel::Kernel::SharedStorage), sizeof(typename ProjectionKernel::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(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 +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; + 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(mat->src[1]->data); + const auto* w = static_cast(mat->src[0]->data); + const auto* xs = reinterpret_cast(x + int64_t(k) * ((int64_t(rows) + 3) / 4 * 4)); + const auto* ws = static_cast(mat->src[2]->data); + const auto* nw = RMS ? static_cast(mul->src[1]->data) : nullptr; + const auto* pe = RMS ? static_cast(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::value ? static_cast(mat->data) : static_cast(output); + if ((reinterpret_cast(x) | reinterpret_cast(w) | reinterpret_cast(xs) | + reinterpret_cast(ws) | reinterpret_cast(nw) | reinterpret_cast(pe) | + reinterpret_cast(address_space)) % + 16 || + reinterpret_cast(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(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(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(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(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(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(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 +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(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(a->data), y = reinterpret_cast(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_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(context.stream); + bool alias = reinterpret_cast(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(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0); + else if (key) + error = projection_launch(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, prefix); + else + error = projection_launch(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(input->data) % 16 == 0 ? prefix_copy : prefix_copy; + kernel<<((count4 + 255) / 256), 256, 0, stream>>>(static_cast(input->data), static_cast(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)); +} diff --git a/src/core/cuda/i8_projection.h b/src/core/cuda/i8_projection.h new file mode 100644 index 0000000000..fb815d0827 --- /dev/null +++ b/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 +#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 diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu new file mode 100644 index 0000000000..45f6ecdf14 --- /dev/null +++ b/src/core/cuda/int8.cu @@ -0,0 +1,739 @@ +#include +#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 +#elif !defined(SD_CUDA_VENDOR_MUSA) +#include +#endif +#include +#include +#include +#include + +namespace SD_CUDA_NAMESPACE { + namespace sd_i8_cuda { + template + 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(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"); + } + static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + int group; + if (!sd_ggml_convrot_describe(node, &group) || 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 && rows <= INT_MAX / (k / 256) && 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; + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(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); + if (quant->src[0] == next && convrot_supports(device, quant, nullptr) && + size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) + return 3; + } + return 2; + } + struct Plan { + ggml_ext_gpu_device_info device; + Kind kind; + 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; + 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 (op.weight_scale && !op.bias && p->k <= 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))) && + (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, large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0); + p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst)); + *workspace = total; + return GGML_STATUS_SUCCESS; + } + static void destroy(void* state, void*) { + auto* p = static_cast(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(a->data), y = reinterpret_cast(b->data); + return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); + } + 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 (shared + 128 <= p.device.shared_memory) { + static std::mutex mutex; + static std::unordered_map configured; + { + std::lock_guard 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<<>>(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<<>>(src, partial, p.k, groups, total_groups); + reduce_convrot_row_amax_i8_cuda<<>>(partial, scales, groups, p.rows); + convrot_group_quantize_i8_cuda<<>>(src, dst, scales, p.k, groups, total_groups, p.k); + } + return cuda_last_error(); + } + template + 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(src)[row * n + col]) * xs[row]; + up = float(reinterpret_cast(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 + 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; + static std::mutex mutex; + static std::unordered_map initialized; + { + std::lock_guard 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<<>>(src, output, + reinterpret_cast(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(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(context.stream); + auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + 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; + void* output = alias ? scratch + p.output : dst->data; + CUDAError error = cuda_success; + if (!mat) { + auto* quant = static_cast(output); + auto* scales = reinterpret_cast(quant + p.k * p.padded); + error = convrot_launch(p, stream, static_cast(dst->src[0]->data), quant, scales, reinterpret_cast(scratch + p.partial)); + } else { + auto* quant = reinterpret_cast(scratch + p.quant); + const float* scales; + if (op.input->type == GGML_TYPE_F32) { + auto* temporary_scales = reinterpret_cast(scratch + p.scales); + scales = temporary_scales; + if (op.group_size) + error = convrot_launch(p, stream, static_cast(op.input->data), quant, temporary_scales, reinterpret_cast(scratch + 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<<>>(static_cast(op.input->data), quant, temporary_scales, p.k, p.rows); + } + } else { + auto* data = static_cast(op.input->data); + scales = op.packed ? reinterpret_cast(data + p.k * p.padded) : static_cast(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(op.weight_scale->data) : nullptr; + const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; + const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0; + auto* gemm_output = reinterpret_cast(static_cast(output) + prefix); + auto* accum = reinterpret_cast(scratch + 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(op.weight->data), scales, ws, + intermediate ? reinterpret_cast(accum) : gemm_output, scratch + p.gemm_offset, stream); + if (status != GGML_STATUS_SUCCESS) + return status; + } else +#endif + { + auto blas_query = reinterpret_cast(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(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(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(accum); + if (p.kind == SWIGLU_QUANTIZE) { + if (p.fast_swiglu) + error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); + else if (p.gemm) + error = quantize_epilogue(p, values, nullptr, nullptr, static_cast(output), stream); + else + error = quantize_epilogue(p, values, scales, ws, static_cast(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<<>>(values, scales, ws, static_cast(output), p.n, p.rows); + else + swiglu<<>>(values, scales, ws, static_cast(output), p.n, p.rows); + error = cuda_last_error(); + } else if (!p.gemm) { + const unsigned gy = unsigned(std::min(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535); + if ((reinterpret_cast(ws) | reinterpret_cast(bias) | reinterpret_cast(gemm_output)) % 16 == 0) { + dequantize_i32_rows_vec4_cuda<<>>( + reinterpret_cast(accum), scales, reinterpret_cast(ws), reinterpret_cast(bias), reinterpret_cast(gemm_output), p.n / 4, p.rows); + } else { + dequantize_i32_rows_cuda<<>>(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)); + } + +} diff --git a/src/core/cuda/modulate.cu b/src/core/cuda/modulate.cu new file mode 100644 index 0000000000..c4af6f6e02 --- /dev/null +++ b/src/core/cuda/modulate.cu @@ -0,0 +1,95 @@ +#include +#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(dst->data); + for (const auto* input : {x, scale, shift}) { + if (input) { + const auto src = reinterpret_cast(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(context.stream); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(e->workspace) : dst->data); + const unsigned blocks = unsigned(std::min((p.count + 255) / 256, 65535)); + modulate<<>>(static_cast(x->data), static_cast(scale->data), + shift ? static_cast(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)); + } +} diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu new file mode 100644 index 0000000000..d75728021f --- /dev/null +++ b/src/core/cuda/registration.cu @@ -0,0 +1,42 @@ +#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); +#endif + +namespace SD_CUDA_NAMESPACE { + void register_rope(ggml_backend_dev_t device); + void register_i8(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_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); +#endif + } +} diff --git a/src/core/cuda/registration.h b/src/core/cuda/registration.h new file mode 100644 index 0000000000..f98c36fb17 --- /dev/null +++ b/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 diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu new file mode 100644 index 0000000000..9f4192b259 --- /dev/null +++ b/src/core/cuda/rms_affine.cu @@ -0,0 +1,266 @@ +#include +#include +#include +#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(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(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 + 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 bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + if (!b) + return false; + const auto x = reinterpret_cast(a->data), y = reinterpret_cast(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*) {} + + template + 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<<>>( + static_cast(op.input->data), static_cast(op.weight->data), + static_cast(op.gate->data), op.residual ? static_cast(op.residual->data) : nullptr, + op.shift ? static_cast(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 = (overlaps(dst, op.input) && dst->data != op.input->data) || + (overlaps(dst, op.residual) && dst->data != op.residual->data) || + overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); + const auto stream = static_cast(context.stream); + CUDAError error = cuda_invalid_value; + switch (op.kind) { + case Kind::Residual: + error = launch(stream, op, output); + break; + case Kind::SelfResidual: + error = launch(stream, op, output); + break; + case Kind::Modulate: + error = launch(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(output); + const auto error = op.kind == sd_rms_affine::Kind::SelfResidual + ? sd_rms_affine::launch(stream, op, dst) + : sd_rms_affine::launch(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)); + } +} diff --git a/src/core/cuda/rms_affine.h b/src/core/cuda/rms_affine.h new file mode 100644 index 0000000000..2bd1abcdce --- /dev/null +++ b/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 diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu new file mode 100644 index 0000000000..99cc824f1f --- /dev/null +++ b/src/core/cuda/rope.cu @@ -0,0 +1,344 @@ +#include "core/cuda/rope.h" +#include "core/ggml_extensions.h" + +#include +#include +#include +#include + +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 + 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::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 + 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 + 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<<((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \ + static_cast(x->data), static_cast(nodes[1]->src[1]->data), \ + static_cast(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(stream, nodes, static_cast(output), scale) + : rms_rope_launch(stream, nodes, static_cast(output), scale); + } + return rope_interleaved(nodes[2]) ? rms_rope_launch(stream, nodes, static_cast(output), scale) + : rms_rope_launch(stream, nodes, static_cast(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(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*) { + *workspace = ggml_nbytes(nodes[count - 1]); + 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(a->data), y = reinterpret_cast(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(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 : rope_apply_f32; + kernel<<((pairs + 255) / 256), 256, 0, stream>>>(static_cast(x->data), static_cast(pe->data), + static_cast(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)); + } + +} diff --git a/src/core/cuda/rope.h b/src/core/cuda/rope.h new file mode 100644 index 0000000000..0d45f689e3 --- /dev/null +++ b/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 diff --git a/src/core/cuda/runtime.h b/src/core/cuda/runtime.h new file mode 100644 index 0000000000..f0db1dcbfa --- /dev/null +++ b/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 +#include +#define SD_CUDA_NAMESPACE sd_hip +#define SD_CUDA_NAME "hip" +#elif defined(SD_CUDA_VENDOR_MUSA) +#include +#include +#define SD_CUDA_NAMESPACE sd_musa +#define SD_CUDA_NAME "musa" +#else +#include +#include +#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 + 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 + 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 + static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) { + return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes); + } +#endif + template + 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 + 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_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_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 diff --git a/src/core/cuda/sage_attention.cu b/src/core/cuda/sage_attention.cu new file mode 100644 index 0000000000..bb42d5b9df --- /dev/null +++ b/src/core/cuda/sage_attention.cu @@ -0,0 +1,387 @@ +#include "core/ggml_extensions.h" +#include "runtime.h" + +#if __CUDACC_VER_MAJOR__ >= 12 + +#include +#include +#include +#include +#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 + 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 + 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 + 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 + 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 + 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 + 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(data + work.q); + auto* ki = reinterpret_cast(data + work.k); + auto* vi = reinterpret_cast(data + work.v); + auto* qs = reinterpret_cast(data + work.q_scale); + auto* ks = reinterpret_cast(data + work.k_scale); + auto* vs = reinterpret_cast(data + work.v_scale); + auto* km = reinterpret_cast(data + work.k_mean); + auto* out = reinterpret_cast(data + work.output); + sage_key_mean<<>>(static_cast(k->data), km, lk); + sage_quantize_per_thread<<>>( + static_cast(q->data), nullptr, qi, qs, lq); + sage_quantize_per_thread<<>>( + static_cast(k->data), km, ki, ks, lk); + // The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++. + sage_value_scale<<>>(static_cast(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f); + sage_quantize_value<<>>(static_cast(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<<>>( + 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(dst->data), ggml_nelements(dst)); + } + + template + 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(stream, dst, data, work, scale); + return; + } + if (mode == GGML_SAGE_ATTN_2) { + sage_attn_fp8(stream, dst, data, work, scale); + return; + } + auto* qi = reinterpret_cast(data + work.q); + auto* ki = reinterpret_cast(data + work.k); + auto* qs = reinterpret_cast(data + work.q_scale); + auto* ks = reinterpret_cast(data + work.k_scale); + auto* km = reinterpret_cast(data + work.k_mean); + auto* out = reinterpret_cast(data + work.output); + + sage_key_mean<<>>(static_cast(k->data), km, lk); + sage_quantize<<>>( + static_cast(q->data), nullptr, qi, qs, lq); + sage_quantize<<>>( + static_cast(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<<>>( + qi, ki, static_cast(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(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(user_data); + SDSageAttentionParams params; + if (!sd_ggml_sage_attn_describe(op, ¶ms) || 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(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], ¶ms)) + return GGML_STATUS_FAILED; + const auto mode = sage_mode(*static_cast(user_data), params.mode); + auto* data = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + auto stream = static_cast(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 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 diff --git a/src/core/cuda/sageattention/LICENSE b/src/core/cuda/sageattention/LICENSE new file mode 100644 index 0000000000..927221d60e --- /dev/null +++ b/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. diff --git a/src/core/cuda/sageattention/README.md b/src/core/cuda/sageattention/README.md new file mode 100644 index 0000000000..ea43a31293 --- /dev/null +++ b/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 `` 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. diff --git a/src/core/cuda/sageattention/attn.cuh b/src/core/cuda/sageattention/attn.cuh new file mode 100644 index 0000000000..540d582273 --- /dev/null +++ b/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 + __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::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float"); + static_assert(std::is_same::value || std::is_same::value, "DTypeSVAccum must be float or half"); + static_assert(std::is_same::value || std::is_same::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::value, "fuse_v_mean only supports half"); + static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation + + using DTypeOut2 = typename std::conditional::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(); + } 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() * 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(); + } 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() * 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::value) { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) { + RO[fq][fv][k] = 0.0f; + } + } else if constexpr (std::is_same::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 smem_Q(smem); + smem_t 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 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 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() * WARP_Q + lane_id % 16, lane_id / 16); + uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_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() * WARP_K + lane_id % 16, lane_id / 16); + + // for causal masking + uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q() * WARP_Q + lane_id / 4; + uint32_t K_idx_lane_base = get_warp_idx_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( + &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( + &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( + &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( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + 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::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d(RS_f16, d); + } + + __syncthreads(); + + // load K + load_global_to_share( + &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( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } else { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + // load V + load_global_to_share( + &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( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + 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(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d(RS_f16, d); + } + + __syncthreads(); + + // load K with predicate + load_global_to_share( + &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( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } else { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + // load V with predicate + load_global_to_share( + &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( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + 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(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // check out of bound in the last iter + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f16[num_tiles_q][num_tiles_k][4]; + RS_32_to_16(RS_f32, RS_f16); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d(RS_f16, d); + } + + // ensure V is ready + cp_async::wait_group<0>(); + __syncthreads(); + + if constexpr (!use_inst_buffer) { + compute_fp16_sv_permuted( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } else { + compute_fp16_sv_permuted_inst_buf( + smem_V, RS_f16, RO, d, V_smem_offset_mma); + } + + __syncthreads(); + } + + // TODO: thread block sync mdo state for num_warps_k > 0 + + normalize_d(RO, m, d); + + // save the result + // if (get_warp_idx_k() == 0) + // { + + // convert half to bfloat16 + if constexpr (std::is_same::value && std::is_same::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() * 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::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::value) { + ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]); + } else if constexpr (std::is_same::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::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() + 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() * 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(offset_O); + } + + offset_O = smem_O.advance_offset_by_row(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(); + 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 diff --git a/src/core/cuda/sageattention/attn_fp8.cuh b/src/core/cuda/sageattention/attn_fp8.cuh new file mode 100644 index 0000000000..838e4c6636 --- /dev/null +++ b/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 + __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::value, "DTypeSVAccum must be float, half is WIP"); + static_assert(std::is_same::value || std::is_same::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(); + } 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() * 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(); + } 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() * 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::value) { +#pragma unroll + for (uint32_t k = 0; k < 8; k++) { + RO[fq][fv][k] = 0.0f; + } + } else if constexpr (std::is_same::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 smem_Q(smem); + smem_t 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 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 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() * WARP_Q + lane_id % 16, lane_id / 16); + uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_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() * 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() * 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() * WARP_Q + lane_id / 4; + uint32_t K_idx_lane_base = get_warp_idx_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( + &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( + &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( + &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( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + 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::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d_f8(RS_f8, d); + } + + __syncthreads(); + + // load K without predicate + load_global_to_share( + &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( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } else { + if constexpr (!use_pv_fp16_accu) { + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } else { + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + __syncthreads(); + // load V + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &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( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + 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(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + // apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d_f8(RS_f8, d); + } + + __syncthreads(); + + // load K with predicate + load_global_to_share( + &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( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } else { + if constexpr (!use_pv_fp16_accu) { + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } else { + compute_fp8_sv_inst_buf_fp16_accu( + smem_V, RS_f8, RO, d); + } + } + + __syncthreads(); + // load V + // for fp16: + // load_global_to_share stride_seq_v + load_fp8_V_global_to_share( + &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( + smem_K, RS, RQ, K_smem_offset_mma); + } else { + compute_int_qk( + 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(Q_idx_lane_base, K_idx_lane_base, RS_f32); + } + apply_out_of_bound_mask(K_idx_lane_base, RS_f32, kv_len); + K_idx_lane_base += CTA_K; + + if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } else if constexpr (std::is_same::value) { + update_mdo(RS_f32, RO, m, d, original_sm_scale); + } + + if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) { + accumulate_d(RS_f32, d); + } + + uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4]; + RS_32_to_8(RS_f32, RS_f8); + + if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) { + accumulate_d_f8(RS_f8, d); + } + + // ensure V is ready + cp_async::wait_group<0>(); + __syncthreads(); + + // for fp16: + // compute_fp16_sv_permuted( + // smem_V, RS_f16, RO, d, V_smem_offset_mma); + if constexpr (!use_inst_buffer) { + compute_fp8_sv( + smem_V, RS_f8, RO, d); + } else { + if constexpr (!use_pv_fp16_accu) { + compute_fp8_sv_inst_buf( + smem_V, RS_f8, RO, d); + } else { + compute_fp8_sv_inst_buf_fp16_accu( + 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(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() * 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::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::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::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() + 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() * 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(offset_O); + } + + offset_O = smem_O.advance_offset_by_row(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(); + 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 diff --git a/src/core/cuda/sageattention/attn_utils.cuh b/src/core/cuda/sageattention/attn_utils.cuh new file mode 100644 index 0000000000..45157b577d --- /dev/null +++ b/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 +#include + +#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 + __device__ __forceinline__ uint32_t get_warp_idx_q() { + return get_warp_id() / num_warps_k; + } + + template + __device__ __forceinline__ uint32_t get_warp_idx_k() { + return get_warp_id() % num_warps_k; + } + + template + __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t& smem) { + static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE); + static_assert(std::is_same::value || std::is_same::value); + + constexpr uint32_t pack_size = std::is_same::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(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(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 + __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t& 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::value || std::is_same::value); + + constexpr uint32_t pack_size = std::is_same::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::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(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(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 + __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& 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(smem_offset); + } + + smem_offset = smem.advance_offset_by_row(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 + __device__ __forceinline__ void compute_int_qk(const smem_t& smem_Q, const smem_t& 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(RS[fq][fk], RQ[fq], RK); + } else if constexpr (DTypeQK == DataType::kInt4) { + mma::mma_sync_m16n16k64_row_col_s4s4s32(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(RS[fq][fk], RQ[fq], RK); + } else if constexpr (DTypeQK == DataType::kInt4) { + mma::mma_sync_m16n16k64_row_col_s4s4s32(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 + __device__ __forceinline__ void compute_int_qk(const smem_t& 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(RS[fq][fk], RQ[fq], RK); + } else if constexpr (DTypeQK == DataType::kInt4) { + mma::mma_sync_m16n16k64_row_col_s4s4s32(RS[fq][fk], RQ[fq], RK); + } + } + } + } + + template + __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::value) { + RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + } else if constexpr (std::is_same::value) { + RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); + } + } + } + } + } + + template + __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::value) { + RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]); + } else if constexpr (std::is_same::value) { + RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]); + } + } + } + } + } + + // for DTypeQKAccum float + template + __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::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::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::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 + __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 + __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 + __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 + __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 + __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::value && compute_unit == ComputeUnit::kCudaCore) || + (std::is_same::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 + __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 + __device__ __forceinline__ void compute_fp16_sv(const smem_t& 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_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::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV); + } else if constexpr (std::is_same::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV); + } + } + } + } + } + + template + __device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t& 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::value) { + mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV); + } else if constexpr (std::is_same::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 + __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t& 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::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((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((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 + __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::value) { + RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2]; + } else if constexpr (std::is_same::value) { + RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]); + } + } + } + } + } + + template + __device__ __forceinline__ void compute_fp8_sv(const smem_t& 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() * ((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::value) { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::value) { + // ! Not Implemented + } + } + offset_V = smem_V.advance_offset_by_row<16>(offset_V); + } + } + } + + template + __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t& 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() * ((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::value) { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::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::value) { + mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::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 + __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t& 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() * ((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::value) { + // mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::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::value) { + // mma::mma_sync_m16n16k32_row_col_f8f8f32(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV); + mma::mma_sync_m16n16k32_row_col_f8f8f16(RO_int32[fq][fv], RS_f8[fq][fk], RV); + } else if constexpr (std::is_same::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 diff --git a/src/core/cuda/sageattention/cp_async.cuh b/src/core/cuda/sageattention/cp_async.cuh new file mode 100644 index 0000000000..98ceca5ba0 --- /dev/null +++ b/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 + +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 + __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 + __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) { +#ifdef CP_ASYNC_ENABLED + uint32_t smem_int_ptr = static_cast(__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 + __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(__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 diff --git a/src/core/cuda/sageattention/math.cuh b/src/core/cuda/sageattention/math.cuh new file mode 100644 index 0000000000..b92ac79868 --- /dev/null +++ b/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 +#include + +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 diff --git a/src/core/cuda/sageattention/mma.cuh b/src/core/cuda/sageattention/mma.cuh new file mode 100644 index 0000000000..f473fa7296 --- /dev/null +++ b/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 +#include +#include +#include +#include + +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 + __device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) { +#ifdef LDMATRIX_M8N8X2_ENABLED + uint32_t smem_int_ptr = static_cast(__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 + __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) { +#ifdef LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__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 + __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) { +#ifdef LDMATRIX_M8N8X4_ENABLED + uint32_t smem_int_ptr = static_cast(__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 + __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 + __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 + __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 + __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 + __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 + __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 + __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 + __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 + __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 + __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(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 + __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 diff --git a/src/core/cuda/sageattention/numeric_conversion.cuh b/src/core/cuda/sageattention/numeric_conversion.cuh new file mode 100644 index 0000000000..8da67d5f13 --- /dev/null +++ b/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 +#include +#include +#include +#include + +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(dst); + } + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sageattention/permuted_smem.cuh b/src/core/cuda/sageattention/permuted_smem.cuh new file mode 100644 index 0000000000..37ad770402 --- /dev/null +++ b/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 +#include +#include + +#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 + 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 + __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 + __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 + 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 + 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 + 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 + __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( + smem_ptr, reinterpret_cast(gptr), predicate); + } + + template + __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const { + b128_t* smem_ptr = base + offset; + cp_async::load_128b(smem_ptr, reinterpret_cast(gptr)); + } + + template + __device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const { + *reinterpret_cast(gptr) = *(base + offset); + } + }; + +} // namespace sageattention + +#endif diff --git a/src/core/cuda/sol_attention.cu b/src/core/cuda/sol_attention.cu new file mode 100644 index 0000000000..0b580abd98 --- /dev/null +++ b/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 +#include +#include + +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(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(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(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(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 + 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(probabilities + d * 16), SOL_BLOCK / 2); + sol_load(vt, reinterpret_cast(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(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(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(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, ¶ms)) + 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(tensor->data) & 3) + return false; + return !(reinterpret_cast(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, ¶ms)) + return GGML_STATUS_FAILED; + const auto* q = dst->src[0]; + const SolWorkspace work(dst); + char* data = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + auto* qb = reinterpret_cast(data + work.q); + auto* kb = reinterpret_cast(data + work.k); + auto* vb = reinterpret_cast(data + work.v); + auto* kc = reinterpret_cast(data + work.kc); + auto* vc = reinterpret_cast(data + work.vc); + auto* qm = reinterpret_cast(data + work.qm); + auto* stats = reinterpret_cast(data + work.stats); + auto* threshold = reinterpret_cast(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(context.stream); + sol_pack<<>>( + static_cast(q->data), static_cast(dst->src[1]->data), static_cast(dst->src[2]->data), + qb, kb, vb, length, work.tokens); + sol_summaries<<>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks); + sol_statistics<<>>(kc, stats, blocks, work.blocks); + sol_thresholds<<>>(qm, stats, threshold, blocks, scale, params.tau); + sol_forward<<>>(qb, kb, vb, kc, vc, threshold, + static_cast(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 diff --git a/src/core/ggml_compat.cpp b/src/core/ggml_compat.cpp new file mode 100644 index 0000000000..061b337067 --- /dev/null +++ b/src/core/ggml_compat.cpp @@ -0,0 +1,185 @@ +#include "core/ggml_compat.h" +#include "core/ggml_extend_backend.h" + +#ifdef SD_GGML_GRAPH_COMPAT +// Upstream ggml still requires its private graph layout for these operations. +#include "ggml-impl.h" + +#include +#include +#include +#include + +int ggml_graph_n_leafs(const ggml_cgraph* graph) { + return graph->n_leafs; +} + +ggml_tensor* ggml_graph_leaf(const ggml_cgraph* graph, int index) { + if (index < 0) { + index += graph->n_leafs; + } + GGML_ASSERT(index >= 0 && index < graph->n_leafs); + return graph->leafs[index]; +} + +void ggml_graph_add_leaf(ggml_cgraph* graph, ggml_tensor* tensor) { + GGML_ASSERT(tensor && graph->n_leafs < graph->size); + graph->leafs[graph->n_leafs++] = tensor; +} + +int32_t ggml_graph_use_count(const ggml_cgraph* graph, const ggml_tensor* tensor) { + if (!graph || !tensor || !graph->visited_hash_set.size) { + return 0; + } + const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor); + return index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index) + ? graph->use_counts[index] + : 0; +} + +bool ggml_graph_cpy_mapped(const ggml_cgraph* src, ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data) { + if (!src || !dst || src == dst || !map || dst->grads || dst->grad_accs || + src->n_nodes > dst->size || src->n_leafs > dst->size) { + return false; + } + std::unordered_map mapping; + std::unordered_set mapped; + std::vector sources; + bool valid = true; + auto remap = [&](const ggml_tensor* source) { + const auto found = mapping.find(source); + if (found != mapping.end()) { + return found->second; + } + auto* tensor = map(source, user_data); + if (tensor && !mapped.insert(tensor).second) { + valid = false; + } + mapping.emplace(source, tensor); + sources.push_back(source); + return tensor; + }; + std::vector nodes(src->n_nodes), leafs(src->n_leafs); + for (int i = 0; i < src->n_nodes; ++i) { + if (!(nodes[i] = remap(src->nodes[i])) || !valid) { + return false; + } + } + for (int i = 0; i < src->n_leafs; ++i) { + if (!(leafs[i] = remap(src->leafs[i])) || !valid) { + return false; + } + } + std::vector> counts; + for (size_t i = 0; i < sources.size(); ++i) { + const auto* source = sources[i]; + if (source->view_src) { + remap(source->view_src); + } + for (const auto* input : source->src) { + if (input) { + remap(input); + } + } + if (!valid) { + return false; + } + auto* tensor = mapping.at(source); + if (!tensor || !src->visited_hash_set.size) { + continue; + } + const size_t index = ggml_hash_find(&src->visited_hash_set, source); + if (index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, index)) { + continue; + } + if (counts.size() >= dst->visited_hash_set.size) { + return false; + } + counts.emplace_back(tensor, src->use_counts[index]); + } + ggml_graph_clear(dst); + dst->order = src->order; + dst->uid = 0; + for (auto* tensor : nodes) { + ggml_graph_add_node(dst, tensor); + } + for (auto* tensor : leafs) { + ggml_graph_add_leaf(dst, tensor); + } + for (const auto& entry : counts) { + const size_t index = ggml_hash_insert(&dst->visited_hash_set, entry.first); + dst->use_counts[index] = entry.second; + } + return true; +} + +bool ggml_graph_cpy_use_counts(const ggml_cgraph* src, ggml_cgraph* dst) { + if (!src || !dst || !src->visited_hash_set.size || !dst->visited_hash_set.size) { + return false; + } + std::unordered_set represented; + for (int i = 0; i < dst->n_nodes; ++i) { + represented.insert(dst->nodes[i]); + } + for (int i = 0; i < dst->n_leafs; ++i) { + represented.insert(dst->leafs[i]); + } + size_t occupied = 0; + for (size_t i = 0; i < dst->visited_hash_set.size; ++i) { + occupied += ggml_bitset_get(dst->visited_hash_set.used, i); + } + std::vector> counts; + for (auto* tensor : represented) { + if (!tensor) { + return false; + } + const size_t source_index = ggml_hash_find(&src->visited_hash_set, tensor); + if (source_index == GGML_HASHSET_FULL || !ggml_bitset_get(src->visited_hash_set.used, source_index)) { + continue; + } + const size_t index = ggml_hash_find(&dst->visited_hash_set, tensor); + if (index == GGML_HASHSET_FULL || !ggml_bitset_get(dst->visited_hash_set.used, index)) { + if (++occupied > dst->visited_hash_set.size) { + return false; + } + } + counts.emplace_back(tensor, src->use_counts[source_index]); + } + for (const auto& entry : counts) { + const size_t index = ggml_hash_find_or_insert(&dst->visited_hash_set, entry.first); + dst->use_counts[index] = entry.second; + } + return true; +} + +ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend, ggml_cgraph* graph, int first, int last) { + if (!backend || !graph || first < 0 || last < first || last > graph->n_nodes) { + return GGML_STATUS_FAILED; + } + if (first == last) { + return GGML_STATUS_SUCCESS; + } + auto view = ggml_graph_view(graph, first, last); + return ggml_backend_graph_compute_async(backend, &view); +} +#endif + +#ifdef SD_USE_UPSTREAM_GGML +#include "ggml-backend-impl.h" +#endif + +ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, void* ptr, size_t size, size_t max_tensor_size) { + auto buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size); +#ifdef SD_USE_UPSTREAM_GGML + if (buffer && !buffer->context) { + auto reg = ggml_backend_dev_backend_reg(device); + if (reg && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) { + // Older Metal backends wrap failed mappings and cannot free their null context. + buffer->iface.free_buffer = nullptr; + ggml_backend_buffer_free(buffer); + return nullptr; + } + } +#endif + return buffer; +} diff --git a/src/core/ggml_compat.h b/src/core/ggml_compat.h new file mode 100644 index 0000000000..9f64262fb8 --- /dev/null +++ b/src/core/ggml_compat.h @@ -0,0 +1,43 @@ +#ifndef __SD_CORE_GGML_COMPAT_H__ +#define __SD_CORE_GGML_COMPAT_H__ + +#include "ggml-backend.h" + +#if __has_include("ggml-graph.h") +#include "ggml-graph.h" +#elif defined(SD_USE_UPSTREAM_GGML) +#define SD_GGML_GRAPH_COMPAT + +int ggml_graph_n_leafs(const struct ggml_cgraph* graph); +struct ggml_tensor* ggml_graph_leaf(const struct ggml_cgraph* graph, int index); +// Like ggml_graph_add_node, appends metadata without traversing tensor dependencies. +void ggml_graph_add_leaf(struct ggml_cgraph* graph, struct ggml_tensor* tensor); +// Returns zero when the tensor has no recorded use count. Includes uses outside graph views. +int32_t ggml_graph_use_count(const struct ggml_cgraph* graph, const struct ggml_tensor* tensor); + +typedef struct ggml_tensor* (*ggml_graph_tensor_map_fn)(const struct ggml_tensor* tensor, void* user_data); + +// Copies forward node/leaf order and recorded use counts, including external consumers. +// The callback must be injective and may return NULL for dependencies outside the +// node/leaf lists. It is called once per tensor in their dependency closure. +// Tensor storage and edges are mapped by the caller. +// The destination must have no gradients and enough capacity. Gradients are not copied. +// Invalid mappings leave the destination unchanged. src and dst must be distinct. +bool ggml_graph_cpy_mapped(const struct ggml_cgraph* src, struct ggml_cgraph* dst, ggml_graph_tensor_map_fn map, void* user_data); + +// Copies recorded counts for tensors in dst's node/leaf lists. Other metadata is unchanged. +// This preserves external consumers when constructing an inference graph segment. +bool ggml_graph_cpy_use_counts(const struct ggml_cgraph* src, struct ggml_cgraph* dst); + +// Executes nodes [first, last) with the original graph's use counts. No graph ownership +// is transferred. Tensor storage must already be allocated; invalid bounds fail. +enum ggml_status ggml_backend_graph_compute_async_range(ggml_backend_t backend, + struct ggml_cgraph* graph, + int first, + int last); + +#else +#error "Patched ggml requires ggml-graph.h; select SD_USE_UPSTREAM_GGML for upstream compatibility" +#endif + +#endif diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 80e2121c6d..90c1008537 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -3,8 +3,10 @@ #include #include #include +#include #include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode) { // reshape A @@ -135,21 +137,32 @@ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_firs // x: [ne3, ne2, ne1, ne0] // return: [ne3, ne2, ne1, ne0/2] - auto x_vec = ggml_ext_chunk(ctx, x, 2, 0, false); - ggml_tensor* gate; - if (gate_first) { - gate = x_vec[0]; - x = x_vec[1]; - } else { - x = x_vec[0]; - gate = x_vec[1]; + GGML_ASSERT(x->ne[0] % 2 == 0); + if (!ggml_is_contiguous_1(x)) { + x = ggml_cont(ctx, x); } - gate = ggml_cont(ctx, gate); - gate = ggml_silu_inplace(ctx, gate); - - x = ggml_mul(ctx, x, gate); // [ne3, ne2, ne1, ne0/2] + return gate_first ? ggml_swiglu(ctx, x) : ggml_swiglu_swapped(ctx, x); +} - return x; +ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate) { +#ifndef SD_USE_UPSTREAM_GGML + if (ctx->backend != nullptr) { + auto out = gate ? ggml_glu_split(ctx->ggml_ctx, gate, hidden, GGML_GLU_OP_SITU) + : ggml_glu(ctx->ggml_ctx, hidden, GGML_GLU_OP_SITU, true); + if (ggml_backend_supports_op(ctx->backend, out)) { + return out; + } + } +#endif + if (!gate) { + const auto dim = hidden->ne[0] / 2; + gate = ggml_ext_slice(ctx->ggml_ctx, hidden, 0, dim, 2 * dim); + hidden = ggml_ext_slice(ctx->ggml_ctx, hidden, 0, 0, dim); + } + auto a = ggml_scale(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, ggml_scale(ctx->ggml_ctx, gate, 0.25f)), 4.f); + a = ggml_mul(ctx->ggml_ctx, a, ggml_sigmoid(ctx->ggml_ctx, gate)); + auto b = ggml_scale(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, ggml_scale(ctx->ggml_ctx, hidden, 1.f / 25.f)), 25.f); + return ggml_mul(ctx->ggml_ctx, a, b); } ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx, @@ -164,6 +177,26 @@ static bool ggml_ext_is_padded_1d(const ggml_tensor* x) { x->nb[3] == x->nb[2] * x->ne[2]; } +ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* shift, + ggml_tensor* scale) { +#ifndef SD_USE_UPSTREAM_GGML + if (x->type == GGML_TYPE_F32 && scale->type == GGML_TYPE_F32 && (!shift || shift->type == GGML_TYPE_F32) && + (sd_backend_is_cpu(ctx->backend) || sd_backend_is(ctx->backend, "CUDA") || + sd_backend_is(ctx->backend, "ROCm") || sd_backend_is(ctx->backend, "MUSA") || sd_backend_is(ctx->backend, "Vulkan"))) { + ggml_tensor* inputs[] = {x, scale, shift}; + return ggml_ext_op_create(ctx->ggml_ctx, sd_ggml_modulate_op(), inputs, shift ? 3 : 2, nullptr); + } +#endif + // Preserve separate multiply/add nodes for rounding and backend fusion matching. + x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, x, scale)); + if (shift != nullptr) { + x = ggml_add(ctx->ggml_ctx, x, shift); + } + return x; +} + ggml_tensor* ggml_ext_scale(ggml_context* ctx, ggml_tensor* x, float factor, @@ -522,21 +555,33 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 1, 3, 2)); x = ggml_reshape_4d(ctx, x, im2col->ne[1], im2col->ne[2], OD, OC * N); } else { - // ggml_conv_3d decomposes into GGML_OP_IM2COL_3D, which some backends - // (e.g. Metal, see #850) do not implement. Fall back to - // GGML_OP_CONV_3D on those backends. - bool im2col_3d_supported = true; + auto convolution = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2); if (backend != nullptr) { - ggml_tensor* im2col = ggml_im2col_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, w->type); - im2col_3d_supported = ggml_backend_supports_op(backend, im2col); - } - if (im2col_3d_supported) { - x = ggml_conv_3d(ctx, w, x, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2); - } else { - int64_t OC = w->ne[3] / IC; - int64_t N = x->ne[3] / IC; - x = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); + // im2col can support large buffers that its reshape/GEMM consumers cannot. + std::vector nodes{convolution}; + bool supported = true; + for (size_t i = 0; i < nodes.size() && supported; ++i) { + auto node = nodes[i]; + if (node == x || node == w) { + continue; + } + supported = ggml_backend_supports_op(backend, node); + for (auto source : node->src) { + if (source != nullptr) { + nodes.push_back(source); + } + } + } + if (!supported) { + int64_t OC = w->ne[3] / IC; + int64_t N = x->ne[3] / IC; + auto fallback = ggml_conv_3d_direct(ctx, w, x, s0, s1, s2, p0, p1, p2, d0, d1, d2, (int)IC, (int)N, (int)OC); + if (ggml_backend_supports_op(backend, fallback)) { + convolution = fallback; + } + } } + x = convolution; } if (scale != 1.f) { @@ -588,9 +633,9 @@ std::vector split_image_qkv(ggml_context* ctx, qkv = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2)); // [3, N, C, H*W] int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 0); // [N, C, H, W] - auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 1); // [N, C, H, W] - auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[3], offset * 2); // [N, C, H, W] + auto q = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 0); // [N, C, H, W] + auto k = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 1); // [N, C, H, W] + auto v = ggml_view_4d(ctx, qkv, W, H, C, N, nb1, nb2, qkv->nb[2], offset * 2); // [N, C, H, W] return {q, k, v}; } @@ -712,21 +757,23 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, // CUDA FA MMA starts at 64 channels; keep the original head's attention scale. q_in = ggml_pad(ctx, q_in, static_cast(64 - d_head), 0, 0, 0); k_in = ggml_pad(ctx, k_in, static_cast(64 - d_head), 0, 0, 0); + v_in = ggml_pad(ctx, v_in, static_cast(64 - d_head), 0, 0, 0); } if (kv_scale != 1.0f) { k_in = ggml_ext_scale(ctx, k_in, kv_scale); } k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16); - v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v_in, 0, 2, 1, 3)); - v_in = ggml_reshape_3d(ctx, v_in, d_head, L_k, n_kv_head * N); - if (pad_head) { - v_in = ggml_pad(ctx, v_in, static_cast(64 - d_head), 0, 0, 0); - } if (kv_scale != 1.0f) { v_in = ggml_ext_scale(ctx, v_in, kv_scale); } - v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); + v_in = ggml_permute(ctx, v_in, 0, 2, 1, 3); + auto v_f16 = ggml_cast(ctx, v_in, GGML_TYPE_F16); + if (!ggml_backend_supports_op(backend, v_f16)) { + v_f16 = ggml_cast(ctx, ggml_ext_cont(ctx, v_in), GGML_TYPE_F16); + } + v_in = v_f16; + v_in = ggml_reshape_3d(ctx, v_in, v_in->ne[0], L_k, n_kv_head * N); if (mask_in != nullptr) { // ggml_flash_attn_ext expects the mask as a contiguous F16 tensor shaped diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 165eb673c9..eda36e2602 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -4,9 +4,12 @@ #include #include +#include "core/ggml_extensions.h" #include "ggml-backend.h" #include "ggml.h" +struct GGMLRunnerContext; + #define EPS 1e-05f static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160"); @@ -49,6 +52,9 @@ std::vector ggml_ext_chunk(ggml_context* ctx, ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true); +// Without gate, hidden contains [hidden, gate] packed along ne[0]. +ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr); + ggml_tensor* ggml_ext_group_norm_32(ggml_context* ctx, ggml_tensor* a); @@ -57,6 +63,13 @@ ggml_tensor* ggml_ext_scale(ggml_context* ctx, float factor, bool inplace = false); +// scale and optional shift must already be broadcastable to x in ggml layout. +// Computes (x + x * scale) + shift; F32 uses sd.modulate.v1 when supported by ctx->backend. +ggml_tensor* ggml_ext_modulate(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* shift, + ggml_tensor* scale); + ggml_tensor* ggml_ext_gelu(ggml_context* ctx, ggml_tensor* x, bool inplace = false); diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 9eafe6f1de..e51b98457a 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -1,4 +1,5 @@ #include "core/ggml_extend_backend.h" +#include "core/ggml_extensions.h" #include #include @@ -19,9 +20,8 @@ #include #endif +#include "core/ggml_compat.h" #include "core/util.h" -#include "ggml-backend-impl.h" -#include "ggml-impl.h" #include "stable-diffusion.h" static std::string trim_copy(const std::string& value) { @@ -296,6 +296,10 @@ bool add_rpc_devices(const std::string& servers) { return false; } ggml_backend_reg_t rpc_reg = ggml_backend_reg_by_name("RPC"); + if (!rpc_reg) { + ggml_backend_load_all(); + rpc_reg = ggml_backend_reg_by_name("RPC"); + } if (!rpc_reg) { LOG_ERROR("RPC backend not found, cannot add RPC servers"); return false; @@ -309,12 +313,19 @@ bool add_rpc_devices(const std::string& servers) { for (const auto& server : rpc_servers) { LOG_INFO("Adding RPC server: %s", server.c_str()); auto reg = ggml_backend_rpc_add_server_fn(server.c_str()); - // no return value to check for success but should print errors from the RPC backend if it fails to add the server + if (!reg) { + LOG_ERROR("failed to add RPC server: %s", server.c_str()); + return false; + } ggml_backend_register(reg); } return true; } +bool sd_add_rpc_servers(const char* servers) { + return add_rpc_devices(servers ? servers : ""); +} + static void ggml_backend_load_all_once() { // If the registry already has devices and the CPU backend is present, // assume either static registration or explicit host-side preloading has @@ -538,24 +549,6 @@ ggml_backend_buffer_type_t sd_cpu_repack_buffer_type_for(ggml_backend_dev_t cpu_ return chosen; } -ggml_backend_buffer_t sd_backend_dev_buffer_from_host_ptr(ggml_backend_dev_t device, - void* ptr, - size_t size, - size_t max_tensor_size) { - ggml_backend_buffer_t buffer = ggml_backend_dev_buffer_from_host_ptr(device, ptr, size, max_tensor_size); - if (buffer != nullptr && buffer->context == nullptr) { - ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(device); - if (reg != nullptr && std::strcmp(ggml_backend_reg_name(reg), "Metal") == 0) { - // Metal can wrap a failed mapping in a non-null buffer. Its free callback also - // dereferences the missing context, so only release the outer buffer. - buffer->iface.free_buffer = nullptr; - ggml_backend_buffer_free(buffer); - return nullptr; - } - } - return buffer; -} - static int cuda_compiled_arch(ggml_backend_t backend) { #ifdef SD_USE_CUDA if (!sd_backend_is(backend, "CUDA")) { @@ -684,7 +677,9 @@ bool sd_backend_supports_conv_implicit_gemm(ggml_backend_t backend, bool conv3d) ggml_backend_t sd_backend_cpu_init() { ggml_backend_load_all_once(); - return ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + auto backend = ggml_backend_init_by_type(GGML_BACKEND_DEVICE_TYPE_CPU, nullptr); + sd_ggml_register_extensions(backend); + return backend; } ggml_backend_t sd_backend_blas_init() { @@ -730,21 +725,30 @@ bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) { return false; } -static ggml_cgraph sd_ggml_graph_view(ggml_cgraph* cgraph0, int i0, int i1) { - ggml_cgraph cgraph = { - /*.size =*/0, - /*.n_nodes =*/i1 - i0, - /*.n_leafs =*/0, - /*.nodes =*/cgraph0->nodes + i0, - /*.grads =*/nullptr, - /*.grad_accs =*/nullptr, - /*.leafs =*/nullptr, - /*.use_counts =*/cgraph0->use_counts, - /*.visited_hash_set =*/cgraph0->visited_hash_set, - /*.order =*/cgraph0->order, - /*.uid =*/0, - }; - return cgraph; +void sd_backend_set_buffer_cache_enabled(ggml_backend_t backend, bool enabled) { + auto dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr; + if (dev == nullptr) { + return; + } + auto reg = ggml_backend_dev_backend_reg(dev); + using SetBufferCacheEnabled = void (*)(ggml_backend_t, bool); + auto fn = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_backend_set_buffer_cache_enabled")); + if (fn != nullptr) { + fn(backend, enabled); + } +} + +void sd_backend_clear_buffer_cache(ggml_backend_t backend) { + auto dev = backend != nullptr ? ggml_backend_get_device(backend) : nullptr; + if (dev == nullptr) { + return; + } + auto reg = ggml_backend_dev_backend_reg(dev); + using ClearBufferCache = void (*)(ggml_backend_t); + auto fn = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_backend_clear_buffer_cache")); + if (fn != nullptr) { + fn(backend); + } } ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, @@ -769,8 +773,7 @@ ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, need = callback_eval(t, true, callback_eval_user_data); } - ggml_cgraph gv = sd_ggml_graph_view(gf, j0, j1 + 1); - status = ggml_backend_graph_compute_async(backend, &gv); + status = ggml_backend_graph_compute_async_range(backend, gf, j0, j1 + 1); if (status != GGML_STATUS_SUCCESS) { break; } @@ -1296,6 +1299,7 @@ ggml_backend_t SDBackendManager::init_cached_backend(const std::string& name) { } SDBackendHandle handle(backend); + sd_ggml_register_extensions(backend); backends_.emplace(actual_key, std::move(handle)); return backend; } diff --git a/src/core/ggml_extend_backend.h b/src/core/ggml_extend_backend.h index 0be0c49314..80f1181dd1 100644 --- a/src/core/ggml_extend_backend.h +++ b/src/core/ggml_extend_backend.h @@ -107,6 +107,9 @@ ggml_backend_t sd_backend_cpu_init(); ggml_backend_t sd_backend_blas_init(); bool sd_backend_set_n_threads_any(ggml_backend_t backend, int n_threads); bool sd_backend_cpu_set_n_threads(ggml_backend_t backend_cpu, int n_threads); +void sd_backend_set_buffer_cache_enabled(ggml_backend_t backend, bool enabled); +void sd_backend_clear_buffer_cache(ggml_backend_t backend); + ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, ggml_cgraph* gf, sd_graph_eval_callback_t callback_eval, diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp new file mode 100644 index 0000000000..27089f6c3b --- /dev/null +++ b/src/core/ggml_extensions.cpp @@ -0,0 +1,441 @@ +#include "core/ggml_extensions.h" + +#ifndef SD_USE_UPSTREAM_GGML +#include "core/cpu/registration.h" +#include "ggml-ext.h" + +#include +#include +#include +#include +#include +#include + +#ifdef SD_BACKEND_EXTENSIONS +#include "sd_backend_extensions.h" +#endif + +static bool modulate_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) { + if (count != 2 && count != 3) { + return false; + } + const auto* x = inputs[0]; + for (int i = 0; i < count; ++i) { + const auto* t = inputs[i]; + if (!t || t->type != GGML_TYPE_F32) { + return false; + } + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (t->ne[d] <= 0 || x->ne[d] % t->ne[d] != 0 || t->nb[d] % sizeof(float)) { + return false; + } + } + } + out->type = GGML_TYPE_F32; + std::memcpy(out->ne, x->ne, sizeof(out->ne)); + return true; +} + +ggml_ext_op_t sd_ggml_modulate_op() { + static const auto op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.modulate.v1", 0, modulate_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_modulate_describe(const ggml_tensor* node) { + if (!node || ggml_ext_op_get(node) != sd_ggml_modulate_op() || !ggml_is_contiguous(node) || node->src[3]) { + return false; + } + ggml_ext_tensor_desc out; + return modulate_infer(node->src, node->src[2] ? 3 : 2, nullptr, &out) && + node->type == out.type && std::memcmp(node->ne, out.ne, sizeof(out.ne)) == 0; +} + +size_t sd_ggml_modulate_workspace_size(const ggml_tensor* node) { + // Owning custom outputs cannot reuse live input storage. Segment outputs may + // already be allocated, so reserve scratch only if their ranges can overlap. + if (!node->data && !node->view_src) { + return 0; + } + if (node->view_src || !node->buffer) { + return ggml_nbytes(node); + } + for (const auto* input : node->src) { + if (!input) { + continue; + } + const size_t bytes = ggml_nbytes(input); + size_t offset = 0; + while (!input->data && input->view_src) { + offset += input->view_offs; + input = input->view_src; + } + if (!input->data) { + continue; + } + if (!input->buffer) { + return ggml_nbytes(node); + } + if (node->buffer != input->buffer) { + continue; + } + const auto dst = reinterpret_cast(node->data); + const auto src = reinterpret_cast(input->data) + offset; + if (dst < src + bytes && src < dst + ggml_nbytes(node)) { + return ggml_nbytes(node); + } + } + return 0; +} + +static bool rope_infer(ggml_tensor* const* inputs, int count, const void*, ggml_ext_tensor_desc* out) { + if (count != 2) { + return false; + } + const auto* x = inputs[0]; + const auto* pe = inputs[1]; + if (x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || x->ne[0] % 2 != 0 || + 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]) { + return false; + } + out->type = GGML_TYPE_F32; + out->ne[0] = x->ne[0]; + out->ne[1] = x->ne[2]; + out->ne[2] = x->ne[1] * x->ne[3]; + out->ne[3] = 1; + return true; +} + +ggml_ext_op_t sd_ggml_rope_op() { + static const ggml_ext_op_t op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.rope_apply.v1", sizeof(int32_t), rope_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_rope_describe(const ggml_tensor* node) { + if (node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node) || !node->src[0] || !node->src[1]) { + return false; + } + ggml_tensor* inputs[] = {node->src[0], node->src[1]}; + ggml_ext_tensor_desc output; + return rope_infer(inputs, 2, nullptr, &output) && + std::memcmp(output.ne, node->ne, sizeof(output.ne)) == 0; +} + +struct I8Attributes { + int64_t shape[3]; + int32_t group_size; + int32_t flags; +}; +enum I8Flags { I8_PACKED = 1, + I8_INPUT_SCALE = 2, + I8_WEIGHT_SCALE = 4, + I8_BIAS = 8 }; + +static bool valid_group(int group) { + return group == 1 || group == 4 || group == 16 || group == 64 || group == 256; +} +static constexpr int64_t i8_size_limit = int64_t(std::min(uint64_t(INT64_MAX), uint64_t(SIZE_MAX))); +static bool packed_height(int64_t k, int64_t rows, int64_t* height) { + if (k <= 0 || rows <= 0 || rows > i8_size_limit / 4) + return false; + const int64_t padded = (rows + 3) / 4 * 4; + const int64_t scales = (rows * 4 - 1) / k + 1; + if (padded > i8_size_limit / k || scales > i8_size_limit / k - padded) + return false; + *height = padded + scales; + return true; +} +static bool scale_matches(const ggml_tensor* scale, int64_t count) { + return !scale || (scale->type == GGML_TYPE_F32 && ggml_is_contiguous(scale) && ggml_nelements(scale) == count); +} +static bool i8_shape(const SDI8Operation& op, const int64_t* shape) { + const auto* w = op.weight; + const auto* x = op.input; + if (!w || !x || w->type != GGML_TYPE_I8 || w->ne[2] != 1 || w->ne[3] != 1 || + !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || x->ne[0] != w->ne[0] || + (op.group_size != 0 && (!valid_group(op.group_size) || w->ne[0] % op.group_size))) + return false; + int64_t rows = 1; + for (int i = 0; i < 3; ++i) { + if (shape[i] <= 0 || shape[i] > (INT64_MAX - 3) / rows) + return false; + rows *= shape[i]; + } + const int64_t k = w->ne[0], n = w->ne[1]; + int64_t height; + if (n <= 0 || n > (i8_size_limit - 4) / 4 || k <= 0 || k > i8_size_limit - 4 - n * 4 || + rows > (i8_size_limit - 3) / (k + 4 + n * 4) || !packed_height(k, rows, &height)) + return false; + if (op.packed) { + if (x->type != GGML_TYPE_I8 || !valid_group(op.group_size) || op.input_scale || + x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1) + return false; + } else { + if (std::memcmp(shape, x->ne + 1, 3 * sizeof(int64_t))) + return false; + if (x->type == GGML_TYPE_I8) { + if (!op.input_scale || op.group_size != 0) + return false; + } else if (x->type != GGML_TYPE_F32 || op.input_scale) + return false; + } + return scale_matches(op.input_scale, rows) && scale_matches(op.weight_scale, w->ne[1]) && scale_matches(op.bias, w->ne[1]); +} +static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) { + I8Attributes attr; + std::memcpy(&attr, attributes, sizeof(attr)); + if (count < 2 || (attr.flags & ~15)) + return false; + SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & I8_PACKED) != 0}; + int index = 2; + for (auto entry : {std::make_pair(I8_INPUT_SCALE, &op.input_scale), std::make_pair(I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(I8_BIAS, &op.bias)}) { + if (attr.flags & entry.first) { + if (index >= count) + return false; + *entry.second = inputs[index++]; + } + } + if (index != count || !i8_shape(op, attr.shape)) + return false; + output->type = GGML_TYPE_F32; + output->ne[0] = op.weight->ne[1]; + std::memcpy(output->ne + 1, attr.shape, sizeof(attr.shape)); + return true; +} +static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* output) { + int32_t group; + std::memcpy(&group, attributes, sizeof(group)); + if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !valid_group(group) || inputs[0]->ne[0] <= 0 || inputs[0]->ne[0] % group) + return false; + const int64_t k = inputs[0]->ne[0], rows = ggml_nrows(inputs[0]); + int64_t height; + if (!packed_height(k, rows, &height)) + return false; + output->type = GGML_TYPE_I8; + output->ne[0] = k; + output->ne[1] = height; + output->ne[2] = output->ne[3] = 1; + return true; +} +ggml_ext_op_t sd_ggml_i8_op() { + static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(I8Attributes), i8_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }(); + return op; +} +ggml_ext_op_t sd_ggml_convrot_op() { + static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.quantize_i8_convrot.v1", sizeof(int32_t), convrot_infer}; auto result = ggml_ext_register_op(&d); GGML_ASSERT(result); return result; }(); + return op; +} +bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) { + if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node)) + return false; + *op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false}; + if (ggml_ext_op_get(node) == sd_ggml_i8_op()) { + I8Attributes attr; + std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr)); + if (attr.flags & ~15 || std::memcmp(attr.shape, node->ne + 1, sizeof(attr.shape))) + return false; + op->group_size = attr.group_size; + op->packed = (attr.flags & I8_PACKED) != 0; + int index = 2; + if (attr.flags & I8_INPUT_SCALE) + op->input_scale = node->src[index++]; + if (attr.flags & I8_WEIGHT_SCALE) + op->weight_scale = node->src[index++]; + if (attr.flags & I8_BIAS) + op->bias = node->src[index++]; + if (((attr.flags & I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & I8_BIAS) && !op->bias)) + return false; + } else + return false; + return i8_shape(*op, node->ne + 1) && node->ne[0] == op->weight->ne[1]; +} +bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) { + if (!node || !node->src[0] || !ggml_is_contiguous(node)) + return false; + if (ggml_ext_op_get(node) == sd_ggml_convrot_op()) + std::memcpy(group, ggml_ext_op_attributes(node), sizeof(int32_t)); + else + return false; + ggml_ext_tensor_desc output; + return convrot_infer(node->src, 1, group, &output) && output.type == node->type && !std::memcmp(output.ne, node->ne, sizeof(output.ne)); +} +ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group, ggml_tensor* xs) { + if (!w || !x) + return nullptr; + const auto* logical = x; + if (x->type == GGML_TYPE_I8 && !xs) { + int packed_group; + if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group) + return nullptr; + logical = x->src[0]; + } + I8Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0}; + ggml_tensor* inputs[5] = {w, x}; + int count = 2; + if (xs) { + attr.flags |= I8_INPUT_SCALE; + inputs[count++] = xs; + } + if (ws) { + attr.flags |= I8_WEIGHT_SCALE; + inputs[count++] = ws; + } + if (bias) { + attr.flags |= I8_BIAS; + inputs[count++] = bias; + } + return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr); +} +ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) { + int32_t attr = group; + return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr); +} + +static bool sage_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { + if (count != 3 || !attributes) + return false; + SDSageAttentionParams params; + std::memcpy(¶ms, attributes, sizeof(params)); + if (!std::isfinite(params.scale) || params.scale <= 0.f || params.mode < GGML_SAGE_ATTN_AUTO || params.mode > GGML_SAGE_ATTN_2_PLUS_PLUS) + return false; + for (int i = 0; i < count; ++i) { + const auto* t = inputs[i]; + if (!t || !ggml_is_contiguous(t) || t->type != (i == 2 ? GGML_TYPE_F16 : GGML_TYPE_F32)) + return false; + int64_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (t->ne[d] <= 0 || t->ne[d] > INT64_MAX / int64_t(sizeof(float)) / elements) + return false; + elements *= t->ne[d]; + } + } + const auto* q = inputs[0]; + const auto* k = inputs[1]; + if (q->ne[0] != k->ne[0] || !ggml_are_same_shape(k, inputs[2]) || q->ne[2] % k->ne[2] != 0 || q->ne[3] != k->ne[3]) + return false; + *out = {GGML_TYPE_F32, {q->ne[0], q->ne[2], q->ne[1], q->ne[3]}}; + return true; +} + +ggml_ext_op_t sd_ggml_sage_attn_op() { + static const auto op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.sage_attn.v1", sizeof(SDSageAttentionParams), sage_attn_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* params) { + if (!node || ggml_ext_op_get(node) != sd_ggml_sage_attn_op() || node->src[3] || !ggml_is_contiguous(node)) + return false; + ggml_ext_tensor_desc output; + const auto* attributes = ggml_ext_op_attributes(node); + if (!sage_attn_infer(node->src, 3, attributes, &output) || node->type != output.type || std::memcmp(node->ne, output.ne, sizeof(output.ne))) + return false; + std::memcpy(params, attributes, sizeof(*params)); + return true; +} + +ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode) { + ggml_tensor* inputs[] = {q, k, v}; + const SDSageAttentionParams params = {scale, mode}; + return ggml_ext_op_create(ctx, sd_ggml_sage_attn_op(), inputs, 3, ¶ms); +} + +static bool sol_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { + if (count != 3 || !attributes) + return false; + SDSolAttentionParams params; + std::memcpy(¶ms, attributes, sizeof(params)); + if (!std::isfinite(params.scale) || params.scale <= 0.f || !std::isfinite(params.tau)) + return false; + const auto* q = inputs[0]; + if (!q || q->ne[0] != 128) + return false; + for (int i = 0; i < count; ++i) { + const auto* t = inputs[i]; + if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || !ggml_are_same_shape(q, t)) + return false; + } + int64_t elements = 1; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (q->ne[d] <= 0 || q->ne[d] > INT64_MAX / int64_t(sizeof(float)) / elements) + return false; + elements *= q->ne[d]; + } + *out = {GGML_TYPE_F32, {q->ne[0], q->ne[2], q->ne[1], q->ne[3]}}; + return true; +} + +ggml_ext_op_t sd_ggml_sol_attn_op() { + static const auto op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.sol_attn.v1", sizeof(SDSolAttentionParams), sol_attn_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params) { + if (!node || ggml_ext_op_get(node) != sd_ggml_sol_attn_op() || node->src[3] || !ggml_is_contiguous(node)) + return false; + ggml_ext_tensor_desc output; + const auto* attributes = ggml_ext_op_attributes(node); + if (!sol_attn_infer(node->src, 3, attributes, &output) || node->type != output.type || std::memcmp(node->ne, output.ne, sizeof(output.ne))) + return false; + std::memcpy(params, attributes, sizeof(*params)); + return true; +} + +ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau) { + ggml_tensor* inputs[] = {q, k, v}; + const SDSolAttentionParams params = {scale, tau}; + return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, ¶ms); +} + +void sd_ggml_register_extensions(ggml_backend_t backend) { + if (!backend) { + return; + } + auto device = ggml_backend_get_device(backend); + auto reg = ggml_backend_dev_backend_reg(device); + const bool cpu = ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1") != nullptr; + static std::mutex mutex; + static std::unordered_set registered; + std::lock_guard lock(mutex); + if (!registered.insert(device).second) { + return; + } + if (!cpu) { +#ifdef SD_BACKEND_EXTENSIONS + for (auto register_backend : sd_backend_extensions) + register_backend(device); +#endif + return; + } + sd_cpu::register_extensions(device); +} + +ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved) { + sd_ggml_register_extensions(backend); + ggml_tensor* inputs[] = {x, pe}; + int32_t attributes = interleaved ? 1 : 0; + return ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &attributes); +} + +#else +void sd_ggml_register_extensions(ggml_backend_t) {} +#endif diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h new file mode 100644 index 0000000000..86c9121a7a --- /dev/null +++ b/src/core/ggml_extensions.h @@ -0,0 +1,70 @@ +#ifndef __SD_CORE_GGML_EXTENSIONS_H__ +#define __SD_CORE_GGML_EXTENSIONS_H__ + +#include "ggml-backend.h" + +void sd_ggml_register_extensions(ggml_backend_t backend); + +#ifndef SD_USE_UPSTREAM_GGML +#include "ggml-ext.h" + +ggml_ext_op_t sd_ggml_modulate_op(); +bool sd_ggml_modulate_describe(const ggml_tensor* node); +size_t sd_ggml_modulate_workspace_size(const ggml_tensor* node); + +ggml_ext_op_t sd_ggml_rope_op(); +bool sd_ggml_rope_describe(const ggml_tensor* node); +// Apply precomputed rotary matrices, with shapes in ne[] order: +// a: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens]. +// pe stores [[cos, -sin], [sin, cos]] for each pair and token, shared across heads/batches. +// Returns contiguous F32 [head_dim, tokens, heads*batch]. Input strides are preserved. +// Interleaved pairs are (2*i, 2*i+1); otherwise pairs are (i, i+rot_dim/2). +// Channels from rot_dim to head_dim are copied unchanged. +ggml_tensor* sd_ggml_rope_apply(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* x, ggml_tensor* pe, bool interleaved); + +struct SDI8Operation { + const ggml_tensor* weight; + const ggml_tensor* input; + const ggml_tensor* input_scale; + const ggml_tensor* weight_scale; + const ggml_tensor* bias; + int group_size; + bool packed; +}; + +ggml_ext_op_t sd_ggml_i8_op(); +ggml_ext_op_t sd_ggml_convrot_op(); +bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation); +bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size); +ggml_tensor* ggml_mul_mat_i8_tensorwise(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size, ggml_tensor* input_scale = nullptr); +ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size); + +enum ggml_sage_attn_mode { + GGML_SAGE_ATTN_AUTO = 0, + GGML_SAGE_ATTN_2, + GGML_SAGE_ATTN_2_PLUS_PLUS, +}; + +struct SDSageAttentionParams { + float scale; + ggml_sage_attn_mode mode; +}; + +ggml_ext_op_t sd_ggml_sage_attn_op(); +bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* params); +// Q/K: contiguous F32 [D, tokens, heads, batch], V: contiguous F16 in the same layout. +// Output: F32 [D, query heads, query tokens, batch]. CUDA-only, unmasked inference. +ggml_tensor* ggml_sage_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, ggml_sage_attn_mode mode); + +struct SDSolAttentionParams { + float scale; + float tau; +}; + +ggml_ext_op_t sd_ggml_sol_attn_op(); +bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* params); +// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch]. +ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, float scale, float tau); +#endif + +#endif diff --git a/src/core/ggml_graph_cut.cpp b/src/core/ggml_graph_cut.cpp index b10e7a14dc..fc6d8c63cb 100644 --- a/src/core/ggml_graph_cut.cpp +++ b/src/core/ggml_graph_cut.cpp @@ -16,7 +16,7 @@ #include "ggml-alloc.h" #include "ggml-backend.h" -#include "ggml-impl.h" +#include "core/ggml_compat.h" namespace sd::ggml_graph_cut { @@ -35,8 +35,8 @@ namespace sd::ggml_graph_cut { static int graph_leaf_index(ggml_cgraph* gf, const ggml_tensor* tensor) { GGML_ASSERT(gf != nullptr); GGML_ASSERT(tensor != nullptr); - for (int i = 0; i < gf->n_leafs; ++i) { - if (gf->leafs[i] == tensor) { + for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) { + if (ggml_graph_leaf(gf, i) == tensor) { return i; } } @@ -389,7 +389,7 @@ namespace sd::ggml_graph_cut { return fail(segment_label + " references an unavailable cut node " + std::to_string(input.node_index)); } - } else if (input.leaf_index < 0 || input.leaf_index >= gf->n_leafs) { + } else if (input.leaf_index < 0 || input.leaf_index >= ggml_graph_n_leafs(gf)) { return fail(segment_label + " references an invalid leaf " + std::to_string(input.leaf_index)); } @@ -444,15 +444,15 @@ namespace sd::ggml_graph_cut { int leaf_count(ggml_cgraph* gf) { GGML_ASSERT(gf != nullptr); - return gf->n_leafs; + return ggml_graph_n_leafs(gf); } ggml_tensor* leaf_tensor(ggml_cgraph* gf, int leaf_index) { GGML_ASSERT(gf != nullptr); - if (leaf_index < 0 || leaf_index >= gf->n_leafs) { + if (leaf_index < 0 || leaf_index >= ggml_graph_n_leafs(gf)) { return nullptr; } - return gf->leafs[leaf_index]; + return ggml_graph_leaf(gf, leaf_index); } ggml_backend_buffer_t tensor_buffer(const ggml_tensor* tensor) { @@ -548,14 +548,14 @@ namespace sd::ggml_graph_cut { }; std::vector graph_layout(ggml_cgraph* graph, bool include_bindings) { - const size_t graph_size = static_cast(graph->n_leafs) + graph->n_nodes; + const size_t graph_size = static_cast(ggml_graph_n_leafs(graph)) + ggml_graph_n_nodes(graph); GraphLayoutTensors layout_tensors(graph_size); const auto& tensors = layout_tensors.tensors; - for (int i = 0; i < graph->n_leafs; ++i) { - layout_tensors.add(graph->leafs[i]); + for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) { + layout_tensors.add(ggml_graph_leaf(graph, i)); } - for (int i = 0; i < graph->n_nodes; ++i) { - layout_tensors.add(graph->nodes[i]); + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + layout_tensors.add(ggml_graph_node(graph, i)); } for (size_t i = 0; i < tensors.size(); ++i) { layout_tensors.add(tensors[i]->view_src); @@ -567,13 +567,13 @@ namespace sd::ggml_graph_cut { const size_t tensor_fields = 5 + 2 * GGML_MAX_DIMS + GGML_MAX_SRC + GGML_MAX_OP_PARAMS / sizeof(int32_t) + (include_bindings ? 2 : 0); signature.reserve(2 + graph_size + tensors.size() * tensor_fields); - signature.push_back(graph->n_nodes); - signature.push_back(graph->n_leafs); - for (int i = 0; i < graph->n_leafs; ++i) { - signature.push_back(layout_tensors.index(graph->leafs[i])); + signature.push_back(ggml_graph_n_nodes(graph)); + signature.push_back(ggml_graph_n_leafs(graph)); + for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) { + signature.push_back(layout_tensors.index(ggml_graph_leaf(graph, i))); } - for (int i = 0; i < graph->n_nodes; ++i) { - signature.push_back(layout_tensors.index(graph->nodes[i])); + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + signature.push_back(layout_tensors.index(ggml_graph_node(graph, i))); } for (auto tensor : tensors) { signature.push_back(tensor->op); @@ -606,18 +606,18 @@ namespace sd::ggml_graph_cut { const Plan& plan, const std::vector& layout) { GGML_ASSERT(gf != nullptr); - if (plan.leaf_names.size() != static_cast(gf->n_leafs) || + if (plan.leaf_names.size() != static_cast(ggml_graph_n_leafs(gf)) || plan.layout != layout) { return false; } - for (int i = 0; i < gf->n_leafs; ++i) { - if (plan.leaf_names[i] != gf->leafs[i]->name) { + for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) { + if (plan.leaf_names[i] != ggml_graph_leaf(gf, i)->name) { return false; } } size_t cut_index = 0; - for (int i = 0; i < gf->n_nodes; ++i) { - auto node = gf->nodes[i]; + for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) { + auto node = ggml_graph_node(gf, i); if (is_graph_cut_tensor(node)) { if (cut_index >= plan.cut_markers.size() || plan.cut_markers[cut_index].first != i || plan.cut_markers[cut_index].second != node->name) { @@ -654,7 +654,7 @@ namespace sd::ggml_graph_cut { } return ggml_graph_node(gf, input_ref.node_index); } - if (input_ref.leaf_index < 0 || input_ref.leaf_index >= gf->n_leafs) { + if (input_ref.leaf_index < 0 || input_ref.leaf_index >= ggml_graph_n_leafs(gf)) { return nullptr; } return leaf_tensor(gf, input_ref.leaf_index); @@ -735,8 +735,7 @@ namespace sd::ggml_graph_cut { GGML_ASSERT(segment_graph != nullptr); for (ggml_tensor* leaf : leaves) { - GGML_ASSERT(segment_graph->n_leafs < segment_graph->size); - segment_graph->leafs[segment_graph->n_leafs++] = leaf; + ggml_graph_add_leaf(segment_graph, leaf); } for (int output_node_index : segment.output_node_indices) { @@ -753,6 +752,7 @@ namespace sd::ggml_graph_cut { for (int node_idx : segment.internal_node_indices) { ggml_graph_add_node(segment_graph, ggml_graph_node(gf, node_idx)); } + GGML_ASSERT(ggml_graph_cpy_use_counts(gf, segment_graph)); *graph_ctx_out = graph_ctx; return segment_graph; } @@ -888,8 +888,8 @@ namespace sd::ggml_graph_cut { return plan; } plan.layout = graph_layout(gf, false); - for (int i = 0; i < gf->n_leafs; ++i) { - plan.leaf_names.emplace_back(gf->leafs[i]->name); + for (int i = 0; i < ggml_graph_n_leafs(gf); ++i) { + plan.leaf_names.emplace_back(ggml_graph_leaf(gf, i)->name); } plan.compute_buffer_size = measure_graph_compute_buffer(backend, gf, params_tensor_set); diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 106a435870..4e4c8ff778 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -52,8 +52,13 @@ void GGMLRunner::free_params_ctx() { void GGMLRunner::alloc_compute_ctx() { ggml_init_params params; - params.mem_size = static_cast(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead()); - params.mem_buffer = nullptr; + params.mem_size = static_cast(ggml_tensor_overhead() * MAX_GRAPH_SIZE + ggml_graph_overhead()); + if (compute_ctx_buffer_ == nullptr) { + const size_t count = (params.mem_size + sizeof(std::max_align_t) - 1) / sizeof(std::max_align_t); + compute_ctx_buffer_.reset(new std::max_align_t[count]); + } + // Keep tensor addresses stable between steps so CUDA graphs can finish warming up. + params.mem_buffer = compute_ctx_buffer_.get(); params.no_alloc = true; compute_ctx = ggml_init(params); @@ -461,6 +466,10 @@ bool GGMLRunner::runner_start() { if (runner_started_) { return true; } + sd_backend_set_buffer_cache_enabled(runtime_backend, true); + for (auto backend : extra_runtime_backends) { + sd_backend_set_buffer_cache_enabled(backend, true); + } cache_.clear(); workspace_.set_extra_backends(extra_runtime_backends); if (auto manager = residency_manager.lock()) { @@ -483,7 +492,10 @@ void GGMLRunner::runner_end() { } workspace_.release(); workspace_.release_blas_scratch(); + free_compute_ctx(); + compute_ctx_buffer_.reset(); cache_.clear(); + cut_cache_.clear(); logged_compute_bytes_.clear(); logged_segment_count_ = 0; if (auto manager = residency_manager.lock()) { @@ -497,6 +509,10 @@ void GGMLRunner::runner_end() { manager->evict_compute_backend_params(tensors); manager->remove_runtime_owner(reinterpret_cast(this)); } + sd_backend_clear_buffer_cache(runtime_backend); + for (auto backend : extra_runtime_backends) { + sd_backend_clear_buffer_cache(backend); + } runner_started_ = false; } @@ -706,8 +722,11 @@ static size_t add_bytes(size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; } -ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t direct_bytes) { +ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, const std::unordered_set* external_outputs) { auto external_backend = [&](const ggml_tensor* tensor) -> ggml_backend_t { + if (external_outputs != nullptr && external_outputs->count(tensor)) { + return runtime_backend; + } if (!params_tensor_set_.count(tensor)) { return nullptr; } @@ -717,7 +736,7 @@ ComputeWorkspace::Measurement GGMLRunner::measure(ggml_cgraph* graph, size_t dir auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { pin_multi_device_nodes(scheduler, copy, graph); }; - return workspace_.measure(graph, direct_bytes, external_backend, assign_nodes); + return workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device()); } std::vector GGMLRunner::memory_requests( @@ -760,7 +779,7 @@ bool GGMLRunner::fits(const std::vector& requests, return true; } -bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { +bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute) { if (sd_backend_is_cpu(runtime_backend)) { sd_backend_cpu_set_n_threads(runtime_backend, n_threads); } @@ -769,21 +788,29 @@ bool GGMLRunner::execute_segment(ggml_cgraph* graph, int n_threads) { } workspace_.set_accelerator_threads(n_threads); auto scheduler = workspace_.scheduler(); - ggml_status status; - if (scheduler != nullptr) { - // CPU + BLAS keeps every tensor in host memory, so callbacks see the same data as on the plain CPU path. - auto eval_callback = workspace_.uses_blas() ? sd_get_backend_eval_callback() : nullptr; - if (eval_callback == nullptr && sd_get_backend_eval_callback() != nullptr && !multi_device_eval_callback_warned) { - LOG_WARN("%s: eval callback is not supported with the backend scheduler; ignoring", get_desc().c_str()); - multi_device_eval_callback_warned = true; - } - ggml_backend_sched_set_eval_callback(scheduler, eval_callback, sd_get_backend_eval_callback_data()); - status = ggml_backend_sched_graph_compute(scheduler, graph); - } else { - status = sd_backend_graph_compute_with_eval_callback(runtime_backend, graph, - sd_get_backend_eval_callback(), - sd_get_backend_eval_callback_data()); - } + GGML_ASSERT(scheduler != nullptr); + ggml_backend_sched_set_eval_callback(scheduler, sd_get_backend_eval_callback(), sd_get_backend_eval_callback_data()); +#ifndef SD_USE_UPSTREAM_GGML + struct BeforeComputeCallback { + ggml_backend_t backend; + const std::function& callback; + bool invoked = false; + } before_compute{runtime_backend, on_before_compute}; + ggml_backend_sched_set_compute_callback( + scheduler, [](ggml_backend_t backend, void* data) { + auto& hook = *static_cast(data); + if (backend == hook.backend && !hook.invoked) { + hook.invoked = true; + hook.callback(); + } + }, + &before_compute); + const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph); + ggml_backend_sched_set_compute_callback(scheduler, nullptr, nullptr); +#else + on_before_compute(); + const auto status = ggml_backend_sched_graph_compute_async(scheduler, graph); +#endif workspace_.synchronize(); if (status != GGML_STATUS_SUCCESS) { last_compute_status_ = status; @@ -835,27 +862,31 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n if (!assign_graph_cut_layer_split_backends(graph)) { return std::nullopt; } - const auto params = collect_used_param_tensors(graph); - const auto& cached_plan = resolve_graph_cut_plan(graph); - const auto full_measurement = measure(graph, cached_plan.compute_buffer_size); - if (full_measurement.buffers.empty()) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return std::nullopt; - } - auto fits_monolithic = [&]() { - // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits. - constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL; - auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)); - for (auto& request : requests) { - request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom); + const auto params = collect_used_param_tensors(graph); + const auto& cached_plan = resolve_graph_cut_plan(graph); + auto manager = residency_manager.lock(); + const bool can_segment = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && + manager != nullptr && manager->segmented_compute_enabled() && + cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1; + // An empty workspace request is a lower bound on the full graph's memory demand. + bool segmented = can_segment && !fits(memory_requests({{runtime_backend, 0}}, 0), params); + ComputeWorkspace::Measurement full_measurement; + if (!segmented) { + full_measurement = measure(graph); + if (full_measurement.buffers.empty()) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return std::nullopt; } - return fits(requests, params); - }; - auto manager = residency_manager.lock(); - const bool segmented = !is_multi_device() && !sd_backend_is_cpu(runtime_backend) && - manager != nullptr && manager->segmented_compute_enabled() && - cached_plan.valid && cached_plan.has_cuts && cached_plan.segments.size() > 1 && - !fits_monolithic(); + if (can_segment) { + // Planning headroom absorbs allocation estimate drift; execution keeps the normal limits. + constexpr size_t planning_headroom = 128ULL * 1024ULL * 1024ULL; + auto requests = memory_requests(full_measurement.buffers, cache_.pending_bytes(graph)); + for (auto& request : requests) { + request.pending_allocation_bytes = add_bytes(request.pending_allocation_bytes, planning_headroom); + } + segmented = !fits(requests, params); + } + } ggml_graph_cut::Plan monolithic_plan; if (!segmented) { monolithic_plan.segments.emplace_back(); @@ -901,7 +932,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n const bool last = index + 1 == plan.segments.size(); auto fail_segment = [&](const char* phase) { LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", get_desc().c_str(), - index + 1, plan.segments.size(), segment.group_name.c_str(), phase); + index + 1, plan.segments.size(), segment.group_name.c_str(), phase); return std::nullopt; }; cut_cache_.prune(segment.live_cut_names); @@ -929,14 +960,24 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } } segment_cleanup{*this, weights, bindings, segment_context}; - auto measurement = segmented ? measure(segment_graph, segment.compute_buffer_size) : full_measurement; + std::unordered_set cache_outputs; + if (segmented) { + for (int output_index : segment.output_node_indices) { + auto tensor = ggml_graph_node(graph, output_index); + if (ggml_graph_cut::is_graph_cut_tensor(tensor) && segment.future_cut_names.count(tensor->name) && + GraphCutTensorCache::can_bind_output(runtime_backend, tensor)) { + cache_outputs.insert(tensor); + } + } + } + auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement; if (!workspace_.prepare(measurement)) { last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace preparation"); } - const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); - const size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph)); - auto ensure_capacity = [&]() { + const size_t cut_bytes = last ? 0 : cut_cache_.estimate_output_bytes(graph, segment); + size_t new_cache_bytes = add_bytes(cut_bytes, cache_.pending_bytes(segment_graph)); + auto ensure_capacity = [&]() { sync_runtime_residency(); auto requests = memory_requests(measurement.buffers, new_cache_bytes); if (fits(requests, weights.params(index))) { @@ -956,9 +997,6 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n return fail_segment("weight preparation"); } // Preparing weights can execute LoRA graphs and reclaim an idle workspace. - if (!workspace_.measurement_matches(segment_graph, measurement)) { - measurement = measure(segment_graph, segment.compute_buffer_size); - } if (!workspace_.prepare(measurement)) { last_compute_status_ = GGML_STATUS_ALLOC_FAILED; return fail_segment("workspace preparation"); @@ -966,6 +1004,11 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n if (!ensure_capacity()) { return fail_segment("workspace capacity check"); } + if (!cut_cache_.bind_outputs(cache_outputs)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("output cache allocation"); + } + new_cache_bytes = add_bytes(last ? 0 : cut_cache_.estimate_output_bytes(graph, segment), cache_.pending_bytes(segment_graph)); if (!workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { pin_multi_device_nodes(scheduler, current); })) { @@ -988,7 +1031,7 @@ std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n } LOG_DEBUG("%s executing segment %zu/%zu: %s", get_desc().c_str(), index + 1, plan.segments.size(), segment.group_name.c_str()); - if (!execute_segment(segment_graph, n_threads)) { + if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) { return fail_segment("execution"); } auto cache_status = cache_.capture(segment_graph); diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 5b963982a3..bb00ec4bcb 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -141,12 +141,12 @@ struct GGMLRunner { size_t logged_segment_count_ = 0; ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; - sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, size_t direct_bytes); + sd::ComputeWorkspace::Measurement measure(ggml_cgraph* graph, const std::unordered_set* external_outputs = nullptr); std::vector memory_requests(const std::vector& sizes, size_t pending_cache_bytes) const; bool fits(const std::vector& requests, const std::vector& params) const; - bool execute_segment(ggml_cgraph* graph, int n_threads); + bool execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute); std::optional> execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs); protected: @@ -161,8 +161,9 @@ struct GGMLRunner { sd::GraphCutTensorCache cut_cache_; sd::ComputeWorkspace workspace_; ggml_context* compute_ctx = nullptr; - bool runner_started_ = false; - bool graph_active_ = false; + std::unique_ptr compute_ctx_buffer_; + bool runner_started_ = false; + bool graph_active_ = false; size_t max_graph_vram_bytes = 0; bool reuse_device_query_ = false; @@ -170,7 +171,6 @@ struct GGMLRunner { std::vector graph_cut_layer_split_backend_vram_limits_; std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) - bool multi_device_eval_callback_warned = false; std::shared_ptr weight_adapter = nullptr; std::weak_ptr residency_manager; diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index b528f1baf2..256f5db84f 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -219,23 +219,23 @@ namespace sd { for (size_t seg_idx = 0; seg_idx < plan.segments.size(); seg_idx++) { ggml_backend_t backend = backend_by_segment[seg_idx]; const auto& segment = plan.segments[seg_idx]; - for (int node_index : segment.internal_node_indices) { + auto assign_node = [&](int node_index) { if (node_index < 0 || node_index >= n_nodes) { - continue; + return; } ggml_tensor* node = ggml_graph_node(gf, node_index); - if (node != nullptr) { + if (node == nullptr) { + return; + } + if (assignment.node_assignments.find(node) == assignment.node_assignments.end()) { assignment.node_assignments[node] = backend; } + }; + for (int node_index : segment.internal_node_indices) { + assign_node(node_index); } for (int node_index : segment.output_node_indices) { - if (node_index < 0 || node_index >= n_nodes) { - continue; - } - ggml_tensor* node = ggml_graph_node(gf, node_index); - if (node != nullptr) { - assignment.node_assignments[node] = backend; - } + assign_node(node_index); } } diff --git a/src/core/pe_cache.cpp b/src/core/pe_cache.cpp new file mode 100644 index 0000000000..df4aa62cb4 --- /dev/null +++ b/src/core/pe_cache.cpp @@ -0,0 +1,30 @@ +#include "core/pe_cache.h" + +namespace sd { + PECacheKey& PECacheKey::add_shape(const ggml_tensor* tensor) { + add(tensor != nullptr); + if (tensor != nullptr) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + add(tensor->ne[i]); + } + } + return *this; + } + + PECacheKey& PECacheKey::add_shapes(const std::vector& tensors) { + add(tensors.size()); + for (const auto* tensor : tensors) { + add_shape(tensor); + } + return *this; + } + + const Tensor& PECache::get_or_create(const PECacheKey& key, const std::function()>& build) { + auto entry = entries_.find(key.name()); + if (entry != entries_.end() && entry->second.key == key) { + return entry->second.value; + } + auto value = build(); + return entries_.insert_or_assign(key.name(), Entry{key, std::move(value)}).first->second.value; + } +} diff --git a/src/core/pe_cache.h b/src/core/pe_cache.h new file mode 100644 index 0000000000..3539707bdb --- /dev/null +++ b/src/core/pe_cache.h @@ -0,0 +1,82 @@ +#ifndef __SD_CORE_PE_CACHE_H__ +#define __SD_CORE_PE_CACHE_H__ + +#include +#include +#include +#include + +#include "core/tensor.hpp" +#include "ggml.h" + +namespace sd { + class PECacheKey { + std::string name_; + std::string fields_; + + public: + explicit PECacheKey(std::string name) + : name_(std::move(name)) {} + const std::string& name() const { return name_; } + + template || std::is_enum_v, int> = 0> + PECacheKey& add(T value) { + if constexpr (std::is_floating_point_v) { + const double field = static_cast(value); + fields_.push_back('f'); + fields_.append(reinterpret_cast(&field), sizeof(field)); + } else { + const uint64_t field = static_cast(value); + fields_.push_back('i'); + fields_.append(reinterpret_cast(&field), sizeof(field)); + } + return *this; + } + + template + PECacheKey& add(const std::vector& values) { + fields_.push_back('v'); + add(values.size()); + for (const auto& value : values) { + add(value); + } + return *this; + } + + template + PECacheKey& add(const std::set& values) { + return add(std::vector(values.begin(), values.end())); + } + + template + PECacheKey& add(const Tensor& values) { + add(values.shape()); + for (int64_t i = 0; i < values.numel(); ++i) { + add(values.data()[i]); + } + return *this; + } + + PECacheKey& add_shape(const ggml_tensor* tensor); + PECacheKey& add_shapes(const std::vector& tensors); + + bool operator==(const PECacheKey& other) const { + return name_ == other.name_ && fields_ == other.fields_; + } + }; + + class PECache { + struct Entry { + PECacheKey key; + Tensor value; + }; + + std::map entries_; + + public: + const Tensor& get_or_create(const PECacheKey& key, const std::function()>& build); + }; + +} + +#endif // __SD_CORE_PE_CACHE_H__ diff --git a/src/core/runner_cache.cpp b/src/core/runner_cache.cpp index b640bb45a6..9bbcae073d 100644 --- a/src/core/runner_cache.cpp +++ b/src/core/runner_cache.cpp @@ -24,15 +24,7 @@ namespace sd { ggml_free(context); } - std::unique_ptr CachedTensor::copy(ggml_backend_t backend, - const std::string& name, - ggml_tensor* source, - ggml_status& status) { - status = GGML_STATUS_FAILED; - if (ggml_graph_cut::tensor_buffer(source) == nullptr) { - return nullptr; - } - status = GGML_STATUS_ALLOC_FAILED; + std::unique_ptr CachedTensor::allocate(ggml_backend_t backend, const std::string& name, ggml_tensor* source) { auto entry = std::make_unique(); entry->context = ggml_init({2 * ggml_tensor_overhead(), nullptr, true}); if (entry->context == nullptr) { @@ -46,6 +38,22 @@ namespace sd { if (entry->buffer == nullptr) { return nullptr; } + return entry; + } + + std::unique_ptr CachedTensor::copy(ggml_backend_t backend, + const std::string& name, + ggml_tensor* source, + ggml_status& status) { + status = GGML_STATUS_FAILED; + if (ggml_graph_cut::tensor_buffer(source) == nullptr) { + return nullptr; + } + status = GGML_STATUS_ALLOC_FAILED; + auto entry = allocate(backend, name, source); + if (entry == nullptr) { + return nullptr; + } if (source->view_src != nullptr || !ggml_is_contiguous(source) || source->buffer == nullptr) { std::vector data(ggml_nbytes(source)); ggml_backend_tensor_get(source, data.data(), 0, data.size()); @@ -155,6 +163,26 @@ namespace sd { return sd::resident_bytes(tensors_, device); } + bool GraphCutTensorCache::can_bind_output(ggml_backend_t backend, const ggml_tensor* tensor) { + return tensor != nullptr && tensor->view_src == nullptr && tensor->op != GGML_OP_NONE && + ggml_is_contiguous(tensor) && ggml_backend_supports_op(backend, tensor); + } + + bool GraphCutTensorCache::bind_outputs(const std::unordered_set& outputs) { + for (const ggml_tensor* output : outputs) { + auto tensor = const_cast(output); + auto entry = CachedTensor::allocate(backend_, tensor->name, tensor); + if (entry == nullptr) { + return false; + } + tensor->buffer = entry->tensor->buffer; + tensor->data = entry->tensor->data; + tensor->extra = entry->tensor->extra; + tensors_[tensor->name] = std::move(entry); + } + return true; + } + size_t GraphCutTensorCache::estimate_output_bytes( ggml_cgraph* graph, const ggml_graph_cut::Segment& segment) const { @@ -171,6 +199,10 @@ namespace sd { !segment.future_cut_names.count(output->name)) { continue; } + const ggml_tensor* cached = get(output->name); + if (cached != nullptr && output->buffer == cached->buffer && output->data == cached->data) { + continue; + } ggml_tensor* source = ggml_graph_cut::cache_source_tensor(output); const size_t tensor_size = GGML_PAD( ggml_backend_buft_get_alloc_size(buffer_type, source), alignment); @@ -196,6 +228,10 @@ namespace sd { !segment.future_cut_names.count(output->name)) { continue; } + const ggml_tensor* cached = get(output->name); + if (cached != nullptr && output->buffer == cached->buffer && output->data == cached->data) { + continue; + } ggml_status status; auto entry = CachedTensor::copy(backend_, output->name, ggml_graph_cut::cache_source_tensor(output), status); if (entry == nullptr) { diff --git a/src/core/runner_cache.h b/src/core/runner_cache.h index f1b5cf91c1..b6fe1c7027 100644 --- a/src/core/runner_cache.h +++ b/src/core/runner_cache.h @@ -18,6 +18,7 @@ namespace sd { ggml_backend_buffer_t buffer = nullptr; ggml_tensor* tensor = nullptr; ~CachedTensor(); + static std::unique_ptr allocate(ggml_backend_t backend, const std::string& name, ggml_tensor* source); static std::unique_ptr copy(ggml_backend_t backend, const std::string& name, ggml_tensor* source, @@ -59,6 +60,8 @@ namespace sd { size_t resident_bytes(ggml_backend_dev_t device) const; size_t estimate_output_bytes(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment) const; + static bool can_bind_output(ggml_backend_t backend, const ggml_tensor* tensor); + bool bind_outputs(const std::unordered_set& outputs); ggml_status capture(ggml_cgraph* graph, const ggml_graph_cut::Segment& segment, const char* log_desc); void prune(const std::unordered_set& keep_names); void clear() { tensors_.clear(); } diff --git a/src/core/segment_graph_bindings.cpp b/src/core/segment_graph_bindings.cpp index 3c40a3b03e..ba72116f3c 100644 --- a/src/core/segment_graph_bindings.cpp +++ b/src/core/segment_graph_bindings.cpp @@ -112,7 +112,7 @@ namespace sd { input_tensor->data = cached_tensor->data == nullptr ? nullptr : static_cast(static_cast(cached_tensor->data) + - input_tensor->view_offs); + input_tensor->view_offs); input_tensor->extra = cached_tensor->extra; } else { input_tensor->buffer = cached_tensor->buffer; diff --git a/src/core/segment_weight_pipeline.cpp b/src/core/segment_weight_pipeline.cpp index f769794df3..2df9a89046 100644 --- a/src/core/segment_weight_pipeline.cpp +++ b/src/core/segment_weight_pipeline.cpp @@ -178,7 +178,7 @@ namespace sd { disable(); return; } - switch (manager->prefetch_params(owner_id_, params)) { + switch (manager->prepare_prefetch_params(owner_id_, params)) { case WeightPrefetchResult::Scheduled: queued_params_ = std::move(params); queued_segment_ = next_segment; @@ -195,6 +195,14 @@ namespace sd { disable(); } + void SegmentWeightPipeline::submit_next() { + if (queued_segment_ != SIZE_MAX) { + if (auto manager = residency_manager_.lock()) { + manager->submit_prefetched_params(owner_id_); + } + } + } + void SegmentWeightPipeline::clear() { if (auto manager = residency_manager_.lock()) { manager->clear_prefetched_params(owner_id_); diff --git a/src/core/segment_weight_pipeline.h b/src/core/segment_weight_pipeline.h index 598bf9f403..67949ed019 100644 --- a/src/core/segment_weight_pipeline.h +++ b/src/core/segment_weight_pipeline.h @@ -56,6 +56,7 @@ namespace sd { // Prefetch is best effort; segment_start falls back to synchronous loading. void enqueue_next(size_t segment_index, const DeviceMemoryRequest& request); + void submit_next(); }; } diff --git a/src/core/util.cpp b/src/core/util.cpp index c1d463e4da..a691b423c0 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -343,7 +343,7 @@ int32_t sd_get_num_physical_cores() { // enumerate the set of thread siblings, num entries is num cores std::unordered_set siblings; for (uint32_t cpu = 0; cpu < UINT32_MAX; ++cpu) { - std::ifstream thread_siblings("/sys/devices/system/cpu" + std::to_string(cpu) + "/topology/thread_siblings"); + std::ifstream thread_siblings("/sys/devices/system/cpu/cpu" + std::to_string(cpu) + "/topology/thread_siblings"); if (!thread_siblings.is_open()) { break; // no more cpus } @@ -568,50 +568,59 @@ bool parse_strict_bool(const std::string& text, bool& value) { } static std::string build_progress_bar(int step, int steps, char progress_char = '=', bool show_head = true) { - std::string progress = " |"; - int max_progress = 50; - int32_t current = 0; + std::string progress = " |"; + constexpr int max_progress = 50; + int32_t current = 0; if (steps > 0) { - current = (int32_t)(step * 1.f * max_progress / steps); + current = (int32_t)(std::clamp(step, 0, steps) * 1.f * max_progress / steps); } - for (int i = 0; i < 50; i++) { - if (i > current) { - progress += " "; - } else if (show_head && i == current && i != max_progress - 1) { - progress += ">"; - } else { - progress += progress_char; - } - } - progress += "|"; + int head_length = (current == max_progress) ? 0 : 1; + progress += std::string(current, progress_char) + std::string(head_length, show_head ? '>' : progress_char) + std::string(max_progress - current - head_length, ' ') + "|"; return progress; } -static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true) { +static void print_progress_line(int step, int steps, const std::string& speed_text, char progress_char = '=', bool show_head = true, const std::string& eta_text = "") { if (step == 0) { return; } std::string progress = build_progress_bar(step, steps, progress_char, show_head); + int step_width = static_cast(std::to_string(steps).size()); const char* lf = (step == steps ? "\n" : ""); - printf("\r%s %i/%i - %s\033[K%s", progress.c_str(), step, steps, speed_text.c_str(), lf); + printf("\r%s %*i/%i - %10s%s%s\033[K%s", progress.c_str(), step_width, step, steps, speed_text.c_str(), + eta_text.empty() ? "" : " - ETA ", eta_text.c_str(), lf); fflush(stdout); // for linux } +static int show_eta = -1; // -1 - estimate duration, 0 - hide, 1 - show +static float passed_time = 0.f; + void pretty_progress(int step, int steps, float time) { if (sd_progress_cb) { sd_progress_cb(step, steps, time, sd_progress_cb_data); return; } if (step == 0) { + show_eta = -1; // reinitialize for batch generations + passed_time = 0.f; return; } + std::string remainder_text = ""; + if (show_eta && step < steps) { // if timer is not disabled and bar is not filled, + if (show_eta == -1 && step) // estimate the total time once, then + show_eta = (std::abs(time * steps / step) < 15.f) ? 0 : 1; // hide the timer if total is less than 15 seconds + if (show_eta) { + unsigned remainder = (passed_time += time) / step * (steps - step); // average speed; time*(steps-step) for current + remainder_text = remainder < 3600 ? sd_format("%02u:%02u", remainder / 60, remainder % 60) + : sd_format("%u:%02u:%02u", remainder / 3600, (remainder / 60) % 60, remainder % 60); + } + } const char* unit = "s/it"; float speed = time; if (speed < 1.0f && speed > 0.f) { speed = 1.0f / speed; unit = "it/s"; } - print_progress_line(step, steps, sd_format("%.2f%s", speed, unit)); + print_progress_line(step, steps, sd_format("%.2f%s", speed, unit), '=', true, remainder_text); } void pretty_bytes_progress(int step, int steps, uint64_t bytes_processed, float elapsed_seconds) { diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt new file mode 100644 index 0000000000..55acf0bb91 --- /dev/null +++ b/src/core/vulkan/CMakeLists.txt @@ -0,0 +1,71 @@ +find_package(Vulkan COMPONENTS glslc REQUIRED) +set(SD_VULKAN_GENERATED "${CMAKE_CURRENT_BINARY_DIR}/shaders") +file(MAKE_DIRECTORY "${SD_VULKAN_GENERATED}") + +function(sd_vulkan_feature NAME EXTENSION) + set(SOURCE "${SD_VULKAN_GENERATED}/feature_${NAME}.comp") + set(EXTENSIONS "#extension ${EXTENSION} : require") + if(NAME STREQUAL "SD_VK_CM2") + string(PREPEND EXTENSIONS "#extension GL_KHR_cooperative_matrix : require\n") + endif() + file(WRITE "${SOURCE}" "#version 460\n${EXTENSIONS}\nlayout(local_size_x = 1) in;\nvoid main() {}\n") + execute_process(COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=vulkan1.3 "${SOURCE}" + -o "${SD_VULKAN_GENERATED}/feature_${NAME}.spv" + RESULT_VARIABLE RESULT OUTPUT_QUIET ERROR_QUIET) + if(RESULT EQUAL 0) + set(${NAME} ON PARENT_SCOPE) + else() + set(${NAME} OFF PARENT_SCOPE) + message(STATUS "sd.cpp Vulkan: ${EXTENSION} unavailable in glslc") + endif() +endfunction() + +sd_vulkan_feature(SD_VK_DOT GL_EXT_integer_dot_product) +sd_vulkan_feature(SD_VK_CM1 GL_KHR_cooperative_matrix) +sd_vulkan_feature(SD_VK_CM2 GL_NV_cooperative_matrix2) + +function(sd_vulkan_shader NAME SHADER ENABLED) + set(HEADER "${SD_VULKAN_GENERATED}/${NAME}.h") + if(ENABLED) + set(SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${SHADER}.comp") + set(SPV "${SD_VULKAN_GENERATED}/${NAME}.spv") + set(ENV vulkan1.2) + if(NAME MATCHES "_cm2$") + set(ENV vulkan1.3) + endif() + add_custom_command(OUTPUT "${HEADER}" + COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=${ENV} ${ARGN} "${SOURCE}" -o "${SPV}" + COMMAND "${CMAKE_COMMAND}" "-DINPUT=${SPV}" "-DOUTPUT=${HEADER}" "-DNAME=${NAME}" + -P "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake" + DEPENDS "${SOURCE}" "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake" + BYPRODUCTS "${SPV}" VERBATIM) + else() + file(WRITE "${HEADER}" "#include \nnamespace sd_vulkan_shaders { static constexpr uint32_t ${NAME}[] = {0}; }\n") + endif() + set(SD_VULKAN_HEADERS ${SD_VULKAN_HEADERS} "${HEADER}" PARENT_SCOPE) +endfunction() + +# Preserve the native compilation policy: optimization changes RoPE reductions +# and is disabled for KHR cooperative matrix shaders. +sd_vulkan_shader(rope_apply rope_apply ON) +sd_vulkan_shader(modulate modulate ON -O) +sd_vulkan_shader(rms_modulate rms_modulate ON -O) +sd_vulkan_shader(rms_rope_apply rms_rope_apply ON) +sd_vulkan_shader(rms_rope_cast rms_rope_apply ON -DCAST_F16) +sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O) +sd_vulkan_shader(quantize_i8_convrot_subgroup quantize_i8_convrot_subgroup ${SD_VK_DOT} -O) +sd_vulkan_shader(quantize_i8_convrot_register quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSINGLE_PASS=1) +sd_vulkan_shader(swiglu_quantize_i8_convrot quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1) +sd_vulkan_shader(swiglu_quantize_i8_convrot_16384 quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSWIGLU=1 -DSINGLE_PASS=1 -DSTATIC_K=16384) +sd_vulkan_shader(sigmoid_quantize_i8_convrot_6144 quantize_i8_convrot_subgroup ${SD_VK_DOT} -O -DSIGMOID_MUL=1 -DSINGLE_PASS=1 -DSTATIC_K=6144) +sd_vulkan_shader(modulate_convrot modulate_convrot ${SD_VK_DOT} -O) +sd_vulkan_shader(rms_modulate_convrot modulate_convrot ${SD_VK_DOT} -O -DRMS_AFFINE=1) +sd_vulkan_shader(mul_mat_i8_tensorwise mul_mat_i8_tensorwise ${SD_VK_DOT} -O) +sd_vulkan_shader(mul_mat_i8_tensorwise_cm1 mul_mat_i8_tensorwise_cm1 ${SD_VK_CM1}) +sd_vulkan_shader(mul_mat_i8_tensorwise_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O) +sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DSPLIT_SWIGLU=1) +sd_vulkan_shader(mul_mat_i8_residual_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_CM2} -O -DRESIDUAL=1) +add_custom_target(sd-vulkan-shaders DEPENDS ${SD_VULKAN_HEADERS}) +add_dependencies(${SD_LIB} sd-vulkan-shaders) +target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/rms_affine.cpp") +target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}") diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp new file mode 100644 index 0000000000..7e3e67d48d --- /dev/null +++ b/src/core/vulkan/int8.cpp @@ -0,0 +1,802 @@ +#include +#include +#include +#include +#include +#include "core/ggml_extensions.h" +#include "ggml-ext-vulkan.h" +#include "modulate_convrot.h" +#include "mul_mat_i8_residual_cm2.h" +#include "mul_mat_i8_swiglu_split_cm2.h" +#include "mul_mat_i8_tensorwise.h" +#include "mul_mat_i8_tensorwise_cm1.h" +#include "mul_mat_i8_tensorwise_cm2.h" +#include "quantize_i8_convrot.h" +#include "quantize_i8_convrot_register.h" +#include "quantize_i8_convrot_subgroup.h" +#include "rms_modulate_convrot.h" +#include "sigmoid_quantize_i8_convrot_6144.h" +#include "swiglu_quantize_i8_convrot.h" +#include "swiglu_quantize_i8_convrot_16384.h" + +namespace sd_vulkan_i8 { + struct Device { + const ggml_ext_vulkan_api* api; + ggml_ext_vulkan_device_info info; + ggml_ext_vulkan_compute_info compute; + bool subgroup, cm2, cm2_small; + uint32_t cm_m, cm_n, cm_k; + }; + static bool device_info(ggml_backend_dev_t dev, Device& d) { + auto reg = ggml_backend_dev_backend_reg(dev); + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1")); + auto compute = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_compute_v1")); + d = {}; + if (!query || !compute) + return false; + d.api = query(); + auto service = compute(); + if (!d.api->device_info(dev, &d.info) || !service->properties(dev, &d.compute)) + return false; + d.subgroup = d.info.subgroup_size == 32 && d.info.full_subgroups && d.info.subgroup_shuffle && d.compute.subgroup_arithmetic && d.info.max_group_invocations >= 256; + ggml_ext_vulkan_matrix_properties prop; + for (uint32_t i = 0; service->matrix_properties(dev, i, &prop); ++i) { + if (prop.a != GGML_TYPE_I8 || prop.b != GGML_TYPE_I8 || prop.accumulator != GGML_TYPE_I32 || prop.result != GGML_TYPE_I32 || prop.saturating || !prop.m || !prop.n || !prop.k) + continue; + if (!prop.flexible && !prop.workgroup_scope && !d.cm_m && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm1) > 4) { + d.cm_m = prop.m; + d.cm_n = prop.n; + d.cm_k = prop.k; + } + if (prop.flexible && prop.workgroup_scope && prop.invocations == 128 && 128 % prop.m == 0 && 64 % prop.n == 0 && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm2) > 4) { + d.cm2 = true; + d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0; + } + } + return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4; + } + static const ggml_tensor* node(const ggml_cgraph* graph, int i) { + return ggml_graph_node(const_cast(graph), i); + } + static int count_nodes(const ggml_cgraph* graph) { + return ggml_graph_n_nodes(const_cast(graph)); + } + static int param_i32(const ggml_tensor* node, int i) { + return node->op_params[i]; + } + static float param_f32(const ggml_tensor* node, int i) { + float value; + std::memcpy(&value, node->op_params + i, 4); + return value; + } + static bool is_convrot(const ggml_tensor* node) { + int group; + return sd_ggml_convrot_describe(node, &group); + } + static bool is_mat(const ggml_tensor* node) { + SDI8Operation op; + return sd_ggml_i8_describe(node, &op); + } + static bool aligned(const ggml_tensor* t, size_t alignment) { + size_t offset = 0; + while (!t->data && t->view_src) { + offset += t->view_offs; + t = t->view_src; + } + return (reinterpret_cast(t->data) + offset) % alignment == 0; + } + static bool bounded(const Device& d, const ggml_tensor* t, bool descriptor_aligned = true) { + return t && ggml_nbytes(t) <= d.info.max_buffer_range - (d.info.buffer_alignment - 1) && + aligned(t, descriptor_aligned ? d.info.buffer_alignment : 4); + } + static bool can_modulate(const Device& d, int64_t k, bool affine) { + return (k == 4096 || k == 6144 || (affine && k == 5376)) && d.info.max_group_invocations >= 512 && d.info.max_group_size[0] >= 512 && d.compute.shared_memory >= (k + 1024) * 4; + } + static bool can_quantize(const Device& d, const ggml_tensor* dst) { + int group; + if (!sd_ggml_convrot_describe(dst, &group) || group != 256 || !bounded(d, dst) || !bounded(d, dst->src[0])) + return false; + const auto* x = dst->src[0]; + return x->ne[0] <= UINT32_MAX && ggml_nrows(x) <= UINT32_MAX && d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 256; + } + static bool can_matmul(const Device& d, const ggml_tensor* dst) { + SDI8Operation op; + if (!sd_ggml_i8_describe(dst, &op) || !op.packed || op.group_size != 256 || !op.weight_scale || !bounded(d, dst)) + return false; + for (auto* t : {op.weight, op.input, op.weight_scale, op.bias}) + if (t && !bounded(d, t)) + return false; + const uint64_t cols = op.weight->ne[1], rows = ggml_nrows(dst); + return op.weight->ne[0] % 256 == 0 && op.weight->ne[0] <= UINT32_MAX && cols <= UINT32_MAX && rows <= UINT32_MAX && + (cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] && + d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16; + } + static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) { + if (i + 2 >= count_nodes(graph) || !device.subgroup) { + return false; + } + const ggml_tensor* sigmoid = node(graph, i); + const ggml_tensor* mul = node(graph, i + 1); + const ggml_tensor* quant = node(graph, i + 2); + if (sigmoid->op != GGML_OP_UNARY || ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID || + mul->op != GGML_OP_MUL || mul->src[1] != sigmoid || + !is_convrot(quant) || quant->src[0] != mul || + sigmoid->ne[0] != 6144 || !ggml_are_same_shape(sigmoid, mul->src[0]) || + !can_quantize(device, quant)) { + return false; + } + for (const ggml_tensor* tensor : {sigmoid->src[0], mul->src[0]}) { + if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) || + tensor->view_offs % device.info.buffer_alignment != 0) { + return false; + } + } + return true; + } + + static bool can_fuse_swiglu_convrot(const Device& device, const ggml_cgraph* graph, int i) { + if (i + 1 >= count_nodes(graph) || !device.subgroup) { + return false; + } + const ggml_tensor* glu = node(graph, i); + const ggml_tensor* quant = node(graph, i + 1); + if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || + glu->src[1] == nullptr || param_i32(glu, 1) != 0 || + !is_convrot(quant) || quant->src[0] != glu || + !can_quantize(device, quant) || !ggml_is_contiguous(glu->src[0]) || !ggml_is_contiguous(glu->src[1])) { + return false; + } + for (const ggml_tensor* tensor : {glu->src[0], glu->src[1]}) { + if (tensor->type != GGML_TYPE_F32 || tensor->view_offs % device.info.buffer_alignment != 0) { + return false; + } + } + return true; + } + + static int match_modulate_convrot(const Device& device, const ggml_cgraph* graph, int quant_idx, int* indices) { + if (quant_idx < 1 || quant_idx >= count_nodes(graph)) { + return 0; + } + const ggml_tensor* quant = node(graph, quant_idx); + if (!is_convrot(quant) || !can_quantize(device, quant)) { + return 0; + } + const ggml_tensor* nodes[4]{}; + + const ggml_tensor* mod = nullptr; + const bool affine = sd_ggml_modulate_describe(quant->src[0]); + const int count = affine ? 2 : 3; + if (affine) { + nodes[0] = quant->src[0]; + nodes[1] = quant; + mod = nodes[0]->src[1]; + const ggml_tensor* vector = nodes[0]->src[2]; + if (vector && (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) || + vector->ne[0] != nodes[0]->ne[0] || ggml_nelements(vector) != nodes[0]->ne[0])) { + return 0; + } + } else { + const ggml_tensor* mul = quant->src[0]; + if (mul->op != GGML_OP_MUL || mul->src[0]->op != GGML_OP_NORM) { + return 0; + } + nodes[0] = mul->src[0]; + nodes[1] = mul; + nodes[2] = quant; + mod = mul->src[1]; + } + const ggml_tensor* x = nodes[0]->src[0]; + if ((x->ne[0] != 4096 && x->ne[0] != 6144 && x->ne[0] != 5376) || + !can_modulate(device, x->ne[0], affine) || + x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || mod->type != GGML_TYPE_F32 || + !ggml_is_contiguous(mod) || mod->ne[0] != x->ne[0] || ggml_nelements(mod) != x->ne[0]) { + return 0; + } + for (int j = 0; j < count; ++j) { + if (j < count - 1 && (nodes[j]->type != GGML_TYPE_F32 || !ggml_is_contiguous(nodes[j]) || !ggml_are_same_shape(x, nodes[j]))) { + return 0; + } + } + int next = count - 1; + for (int j = quant_idx; j >= std::max(0, quant_idx - 64) && next >= 0; --j) { + if (node(graph, j) == nodes[next]) { + indices[next--] = j; + } + } + return next < 0 ? count : 0; + } + + static int match_rms_modulate_convrot(const Device& device, const ggml_cgraph* graph, int q, int* indices) { + if (!can_modulate(device, 5376, true) || match_modulate_convrot(device, graph, q, indices + 2) != 2) { + return 0; + } + const auto weighted = node(graph, indices[2])->src[0]; + if (weighted->op != GGML_OP_MUL || weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 5376) { + return 0; + } + const auto norm = weighted->src[0]; + const auto x = norm->src[0]; + const auto weight = weighted->src[1]; + if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || weight->type != GGML_TYPE_F32 || + !ggml_is_contiguous(weight) || ggml_nelements(weight) != 5376 || weight->ne[0] != 5376) { + return 0; + } + int next = 1; + for (int j = indices[2] - 1; j >= std::max(0, indices[2] - 64) && next >= 0; --j) { + if (node(graph, j) == (next == 1 ? weighted : norm)) { + indices[next--] = j; + } + } + + return next < 0 ? 4 : 0; + } + + static int can_fuse_modulate_convrot(const Device& device, const ggml_cgraph* graph, int i) { + const ggml_op op = node(graph, i)->op; + if (op == GGML_OP_RMS_NORM) { + int indices[4]; + if (match_rms_modulate_convrot(device, graph, i + 3, indices) != 4) { + return 0; + } + for (int j = 0; j < 4; ++j) { + if (indices[j] != i + j) { + return 0; + } + } + return 4; + } + if (op != GGML_OP_NORM && !sd_ggml_modulate_describe(node(graph, i))) { + return 0; + } + const int count = op == GGML_OP_NORM ? 3 : 2; + int indices[4]; + if (match_modulate_convrot(device, graph, i + count - 1, indices) != count) { + return 0; + } + for (int j = 0; j < count; ++j) { + if (indices[j] != i + j) { + return 0; + } + } + return count; + } + + static bool match_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i, int* indices) { + if (i < 2 || i >= count_nodes(graph) || !device.cm2_small) { + return false; + } + const ggml_tensor* glu = node(graph, i); + if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || + param_i32(glu, 1) != 0 || glu->src[1] == nullptr || glu->ne[0] % 128 != 0 || + glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) || + glu->view_offs % device.info.buffer_alignment != 0) { + return false; + } + const ggml_tensor* gate = glu->src[0]; + const ggml_tensor* up = glu->src[1]; + if (gate == up || gate->src[1] != up->src[1] || !ggml_are_same_shape(gate, up)) { + return false; + } + for (auto mat : {gate, up}) { + if (!is_mat(mat) || mat->view_src != nullptr || mat->src[3] != nullptr || + !can_matmul(device, mat)) { + return false; + } + } + int found = 0; + for (int j = std::max(0, i - 32); j < i; ++j) { + if (node(graph, j) == gate || node(graph, j) == up) { + indices[found++] = j; + } + } + if (found != 2) { + return false; + } + indices[2] = i; + + return true; + } + + static bool can_fuse_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i) { + int indices[3]; + return match_i8_swiglu_split(device, graph, i + 2, indices) && indices[0] == i && indices[1] == i + 1; + } + + static bool can_fuse_i8_swiglu(const Device& device, const ggml_cgraph* graph, int i) { + if (i + 1 >= count_nodes(graph) || !device.cm2_small) { + return false; + } + const ggml_tensor* mat = node(graph, i); + const ggml_tensor* glu = node(graph, i + 1); + return is_mat(mat) && can_matmul(device, mat) && + mat->view_src == nullptr && mat->src[3] == nullptr && mat->ne[0] % 256 == 0 && + glu->op == GGML_OP_GLU && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU && + param_i32(glu, 1) == 0 && glu->src[0] == mat && glu->src[1] == nullptr && + glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) && + glu->view_offs % device.info.buffer_alignment == 0; + } + + static int match_i8_residual(const Device& device, const ggml_cgraph* graph, int add_idx, int* indices) { + if (add_idx < 2 || add_idx >= count_nodes(graph) || !device.cm2) { + return 0; + } + const ggml_tensor* dst = node(graph, add_idx); + if (dst->op != GGML_OP_ADD) { + return 0; + } + const ggml_tensor* mul = dst->src[0]->op == GGML_OP_MUL ? dst->src[0] : dst->src[1]; + if (mul->op != GGML_OP_MUL) { + return 0; + } + const ggml_tensor* scaled_mat = mul->src[0]; + const bool scaled = scaled_mat->op == GGML_OP_SCALE; + const ggml_tensor* mat = scaled ? scaled_mat->src[0] : scaled_mat; + if (!is_mat(mat) || !can_matmul(device, mat) || mat->src[3] != nullptr || + (scaled && param_f32(scaled_mat, 1) != 0.0f)) { + return 0; + } + const int count = scaled ? 4 : 3; + const ggml_tensor* nodes[4] = {mat, scaled ? scaled_mat : mul, scaled ? mul : dst, dst}; + + const ggml_tensor* gate = mul->src[1]; + const ggml_tensor* residual = dst->src[0] == mul ? dst->src[1] : dst->src[0]; + if (!ggml_are_same_shape(mat, dst) || !ggml_are_same_shape(residual, dst) || + gate->ne[0] != mat->ne[0] || ggml_nelements(gate) != mat->ne[0]) { + return 0; + } + for (const ggml_tensor* tensor : {mat, scaled_mat, mul, gate, residual, dst}) { + if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) || + tensor->view_offs % device.info.buffer_alignment != 0) { + return 0; + } + } + int next = count - 1; + for (int j = add_idx; j >= std::max(0, add_idx - 32) && next >= 0; --j) { + if (node(graph, j) == nodes[next]) { + indices[next--] = j; + } + } + return next < 0 ? count : 0; + } + + static int can_fuse_i8_residual(const Device& device, const ggml_cgraph* graph, int i) { + if (i + 2 >= count_nodes(graph) || !is_mat(node(graph, i))) { + return 0; + } + const int count = node(graph, i + 1)->op == GGML_OP_SCALE ? 4 : 3; + int indices[4]; + if (match_i8_residual(device, graph, i + count - 1, indices) != count) { + return 0; + } + for (int j = 0; j < count; ++j) { + if (indices[j] != i + j) { + return 0; + } + } + return count; + } + + static int match_i8_row_concat(const Device& device, const ggml_cgraph* graph, int end, int* indices) { + const ggml_tensor* concat = node(graph, end); + if (concat->op != GGML_OP_CONCAT || param_i32(concat, 0) != 1 || + concat->type != GGML_TYPE_F32 || concat->ne[2] != 1 || concat->ne[3] != 1 || + !ggml_is_contiguous(concat) || !ggml_is_contiguous(concat->src[0]) || + concat->view_offs % device.info.buffer_alignment != 0 || + ggml_nbytes(concat->src[0]) % device.info.buffer_alignment != 0) { + return 0; + } + const ggml_tensor* tail = concat->src[1]; + int count = 0; + for (int j = end - 1; j >= std::max(0, end - 32); --j) { + if (node(graph, j) != tail) { + continue; + } + if (is_mat(tail) && can_matmul(device, tail)) { + indices[0] = j; + count = 1; + } else { + count = match_i8_residual(device, graph, j, indices); + } + break; + } + if (!count) { + return 0; + } + indices[count++] = end; + + return count; + } + + static int can_fuse_i8_row_concat(const Device& device, const ggml_cgraph* graph, int i) { + if (!is_mat(node(graph, i))) { + return 0; + } + for (int end = i + 1; end < std::min(i + 5, count_nodes(graph)); ++end) { + int indices[5]; + const int count = match_i8_row_concat(device, graph, end, indices); + if (count != end - i + 1) { + continue; + } + bool consecutive = true; + for (int j = 0; j < count; ++j) { + consecutive &= indices[j] == i + j; + } + if (consecutive) { + return count; + } + } + return 0; + } + + struct QuantParameters { + uint32_t k, rows; + }; + struct ModulateParameters { + uint32_t rows, x_offset, m_offset, h_offset; + float eps; + uint32_t k, w_offset; + }; + struct MatParameters { + uint32_t n, rows, k, has_bias, row_offset; + float output_scale; + uint32_t group_cols; + }; + static_assert(sizeof(QuantParameters) == 8 && sizeof(ModulateParameters) == 28 && sizeof(MatParameters) == 28, "INT8 push constant layout mismatch"); + template + static ggml_ext_vulkan_pipeline shader(const char* name, const uint32_t (&spirv)[N], uint32_t bindings, uint32_t constants) { + return {name, spirv, sizeof(spirv), bindings, constants, nullptr, 0, 0, false}; + } +#define SD_VK_PIPELINE(NAME, BINDINGS, PARAMS) shader("sd." #NAME ".vulkan", sd_vulkan_shaders::NAME, BINDINGS, sizeof(PARAMS)) + static ggml_ext_vulkan_binding range(ggml_ext_vulkan_binding binding, size_t offset, size_t size) { + binding.offset += offset; + binding.size = size; + return binding; + } + static bool overlap(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) { + return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size; + } + static bool output_binding(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const std::vector& inputs, ggml_ext_vulkan_binding& binding, bool& alias, const ggml_tensor* residual = nullptr) { + binding = {output}; + alias = false; + ggml_ext_vulkan_buffer_info dst, src; + if (!d.api->buffer_info(e->backend, &binding, &dst)) + return false; + for (const auto& input : inputs) { + if (!d.api->buffer_info(e->backend, &input, &src)) + return false; + // The residual epilogue reads each element before writing that same element. + if (residual && input.tensor == residual && dst.storage == src.storage && dst.offset == src.offset && dst.size == src.size) + continue; + alias |= overlap(dst, src); + } + if (alias) + binding = {nullptr, e->workspace, 0, ggml_nbytes(output)}; + return true; + } + static ggml_status finish(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const ggml_ext_vulkan_binding& binding, bool alias, ggml_status status) { + if (status == GGML_STATUS_SUCCESS && alias) { + ggml_ext_vulkan_binding dst = {output}; + return d.api->copy(e->backend, &binding, &dst); + } + return status; + } + static ggml_status quantize(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) { + const auto* output = nodes[count - 1]; + const auto* first = nodes[0]; + const bool modulate = count > 1 && (first->op == GGML_OP_NORM || sd_ggml_modulate_describe(first) || first->op == GGML_OP_RMS_NORM); + const bool rms = modulate && first->op == GGML_OP_RMS_NORM; + const bool affine = modulate && (rms || sd_ggml_modulate_describe(first)); + const bool sigmoid = count > 1 && first->op == GGML_OP_UNARY; + const auto* x = count == 1 ? output->src[0] : first->src[0]; + const auto* up = count == 1 || modulate ? nullptr : sigmoid ? nodes[1]->src[0] + : first->src[1]; + const int extra = rms ? 2 : 0; + const auto* mod = modulate ? nodes[extra + (affine ? 0 : 1)]->src[1] : nullptr; + const auto* shift = modulate ? (affine ? nodes[extra]->src[2] : mod) : nullptr; + const auto* weight = rms ? nodes[1]->src[1] : x; + std::vector inputs = {{x}}; + if (up) + inputs.push_back({up}); + if (modulate) { + inputs.push_back({mod}); + if (shift) + inputs.push_back({shift}); + if (rms) + inputs.push_back({weight}); + } + ggml_ext_vulkan_binding dst; + bool alias; + if (!output_binding(e, d, output, inputs, dst, alias)) + return GGML_STATUS_FAILED; + const uint32_t k = uint32_t(x->ne[0]), rows = uint32_t(ggml_nrows(x)); + const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4); + auto packed = range(dst, 0, scale_offset), scales = range(dst, scale_offset, size_t(rows) * 4); + const uint32_t groups[] = {std::min(rows, d.info.max_group_count[0]), 1, 1}; + ggml_status status; + if (modulate) { + std::vector bindings = {{x}, {mod}, {shift ? shift : mod}, packed, scales}; + if (rms) + bindings.push_back({weight}); + ModulateParameters pc = {rows, 0, 0, 0, affine && !rms ? 0.f : param_f32(first, 0), k, 0}; + uint32_t* offsets[] = {&pc.x_offset, &pc.m_offset, &pc.h_offset, nullptr, nullptr, &pc.w_offset}; + for (size_t i = 0; i < bindings.size(); ++i) + if (offsets[i]) { + ggml_ext_vulkan_buffer_info info; + if (!d.api->buffer_info(e->backend, &bindings[i], &info)) + return GGML_STATUS_FAILED; + *offsets[i] = info.shader_offset / 4; + } + auto pipeline = rms ? SD_VK_PIPELINE(rms_modulate_convrot, 6, ModulateParameters) : SD_VK_PIPELINE(modulate_convrot, 5, ModulateParameters); + const uint32_t constants[] = {k, uint32_t(affine), uint32_t(shift != nullptr)}; + pipeline.constants = constants; + pipeline.constant_count = 3; + status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups); + } else { + ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}}; + auto pipeline = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters); + bool subgroup = d.subgroup && k >= 1024; + uint32_t width = k; + if (subgroup) { + pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters); + if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) { + pipeline = SD_VK_PIPELINE(quantize_i8_convrot_register, 3, QuantParameters); + pipeline.constants = &width; + pipeline.constant_count = 1; + } + } + if (up) { + subgroup = true; + pipeline = k == 16384 ? SD_VK_PIPELINE(swiglu_quantize_i8_convrot_16384, 4, QuantParameters) : SD_VK_PIPELINE(swiglu_quantize_i8_convrot, 4, QuantParameters); + if (sigmoid) + pipeline = SD_VK_PIPELINE(sigmoid_quantize_i8_convrot_6144, 4, QuantParameters); + } + if (subgroup) { + pipeline.full_subgroups = d.info.full_subgroups; + pipeline.required_subgroup_size = d.info.subgroup_size_control ? 32 : 0; + } + QuantParameters pc = {k, rows}; + status = d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups); + } + return finish(e, d, output, dst, alias, status); + } + static uint32_t group_columns(uint32_t cols, uint32_t k, uint32_t tile_cols, bool swiglu) { + if (!cols || !k || !tile_cols) + return 1; + const uint64_t col_tiles = (uint64_t(cols) + tile_cols - 1) / tile_cols; + const uint64_t weight_tile_bytes = uint64_t(tile_cols) * k * (swiglu ? 2 : 1); + const uint64_t growth = std::clamp(uint64_t(k) / cols, uint64_t(1), uint64_t(2)); + return uint32_t(std::clamp(uint64_t(8) * 1024 * 1024 * growth / weight_tile_bytes, uint64_t(1), std::min(col_tiles, uint64_t(32)))); + } + static ggml_status multiply(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) { + const auto* output = nodes[count - 1]; + const auto* concat = output->op == GGML_OP_CONCAT ? output : nullptr; + const auto* tail = concat ? concat->src[1] : output; + const auto* glu = tail->op == GGML_OP_GLU ? tail : nullptr; + const auto* mat = glu ? glu->src[0] : tail; + const ggml_tensor *gate = nullptr, *residual = nullptr; + float output_scale = 1.f; + if (tail->op == GGML_OP_ADD) { + const auto* mul = tail->src[0]->op == GGML_OP_MUL ? tail->src[0] : tail->src[1]; + gate = mul->src[1]; + residual = tail->src[0] == mul ? tail->src[1] : tail->src[0]; + mat = mul->src[0]; + if (mat->op == GGML_OP_SCALE) { + output_scale = param_f32(mat, 0); + mat = mat->src[0]; + } + } + const auto* up_mat = glu ? glu->src[1] : nullptr; + SDI8Operation op, up; + if (!sd_ggml_i8_describe(mat, &op) || (up_mat && !sd_ggml_i8_describe(up_mat, &up))) + return GGML_STATUS_FAILED; + const uint32_t k = uint32_t(op.weight->ne[0]), n = uint32_t(op.weight->ne[1]), rows = uint32_t(ggml_nrows(mat)); + const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4); + std::vector inputs = {{op.weight}, {op.input}, {op.weight_scale}, {op.bias ? op.bias : op.weight_scale}}; + if (gate) { + inputs.push_back({gate}); + inputs.push_back({residual}); + } + if (up_mat) { + inputs.push_back({up.weight}); + inputs.push_back({up.weight_scale}); + } + if (concat && ggml_nbytes(concat->src[0])) + inputs.push_back({concat->src[0]}); + ggml_ext_vulkan_binding dst; + bool alias; + if (!output_binding(e, d, output, inputs, dst, alias, concat ? nullptr : residual)) + return GGML_STATUS_FAILED; + auto destination = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst; + ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}}; + auto dispatch = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows, + std::initializer_list constants, uint32_t row_offset, uint32_t dispatch_rows) { + if (!dispatch_rows) + return GGML_STATUS_SUCCESS; + const uint32_t cols = glu ? uint32_t(glu->ne[0]) : n; + MatParameters pc = {n * (up_mat ? 2u : 1u), rows, k, op.bias ? 1u : 0u, row_offset, output_scale, group_columns(cols, k, tile_cols, glu != nullptr)}; + pipeline.constants = constants.begin(); + pipeline.constant_count = uint32_t(constants.size()); + const uint32_t groups[] = {(cols + tile_cols - 1) / tile_cols, (dispatch_rows + tile_rows - 1) / tile_rows, 1}; + return d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups); + }; + ggml_status status = GGML_STATUS_SUCCESS; + if (d.cm2) { + const bool large_k = k % 4096 == 0 || (rows >= n && k >= n); + const uint32_t full_rows = (glu || n % 128 == 0) ? rows / 128 * 128 : 0; + auto pipeline = residual ? SD_VK_PIPELINE(mul_mat_i8_residual_cm2, 8, MatParameters) : up_mat ? SD_VK_PIPELINE(mul_mat_i8_swiglu_split_cm2, 8, MatParameters) + : SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm2, 6, MatParameters); + const uint32_t tile_cols = glu ? 64 : 128; + status = dispatch(pipeline, tile_cols, 128, {128, 1, glu ? 1u : 0u, tile_cols, glu || large_k ? 128u : 64u}, 0, full_rows); + if (status == GGML_STATUS_SUCCESS) { + const bool small = d.cm2_small && (glu || n % 128 == 0) && rows - full_rows <= 32; + const uint32_t tail_cols = small || glu ? 64 : 128, tail_rows = small ? 32 : glu ? 128 + : 64; + status = dispatch(pipeline, tail_cols, tail_rows, {tail_rows, small ? 2u : 0u, glu ? 1u : 0u, tail_cols, small ? 256u : 64u}, full_rows, rows - full_rows); + } + } else { + uint32_t coop_rows = 0; + if (d.cm_m && k % d.cm_k == 0 && n % (4 * d.cm_n) == 0) { + const uint32_t tile_rows = 4 * d.cm_m; + coop_rows = rows / tile_rows * tile_rows; + auto pipeline = SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm1, 6, MatParameters); + pipeline.full_subgroups = d.info.full_subgroups; + pipeline.required_subgroup_size = d.info.subgroup_size_control ? d.info.subgroup_size : 0; + status = dispatch(pipeline, 4 * d.cm_n, tile_rows, {d.info.subgroup_size, d.cm_m, d.cm_n, d.cm_k}, 0, coop_rows); + } + if (status == GGML_STATUS_SUCCESS) + status = dispatch(SD_VK_PIPELINE(mul_mat_i8_tensorwise, 6, MatParameters), 16, 16, {}, coop_rows, rows - coop_rows); + } + if (status == GGML_STATUS_SUCCESS && concat && ggml_nbytes(concat->src[0])) { + ggml_ext_vulkan_binding src = {concat->src[0]}, prefix = range(dst, 0, ggml_nbytes(concat->src[0])); + status = d.api->copy(e->backend, &src, &prefix); + } + return finish(e, d, output, dst, alias, status); + } + static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) { + Device d; + return device_info(dev, d) && can_quantize(d, n); + } + static bool mat_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) { + Device d; + return device_info(dev, d) && can_matmul(d, n); + } + static int match(ggml_backend_dev_t dev, const ggml_cgraph* graph, int first, void*) { + Device d; + if (!device_info(dev, d) || d.info.fusion_disabled) + return 0; + if (int count = can_fuse_modulate_convrot(d, graph, first)) + return count; + if (int count = can_fuse_i8_row_concat(d, graph, first)) + return count; + if (int count = can_fuse_i8_residual(d, graph, first)) + return count; + if (can_fuse_i8_swiglu_split(d, graph, first)) + return 3; + if (can_fuse_i8_swiglu(d, graph, first)) + return 2; + if (can_fuse_sigmoid_convrot(d, graph, first)) + return 3; + if (can_fuse_swiglu_convrot(d, graph, first)) + return 2; + return 0; + } + static void optimize(ggml_backend_dev_t dev, ggml_cgraph* graph, void*) { + Device d; + if (!device_info(dev, d) || d.info.fusion_disabled || d.compute.graph_optimization_disabled) + return; + for (int last = 1; last < count_nodes(graph); ++last) { + int indices[6]; + int count = match_rms_modulate_convrot(d, graph, last, indices); + if (!count) + count = match_modulate_convrot(d, graph, last, indices); + if (!count) + count = match_i8_row_concat(d, graph, last, indices); + if (!count) + count = match_i8_residual(d, graph, last, indices); + if (!count && match_i8_swiglu_split(d, graph, last, indices)) + count = 3; + if (!count) + continue; + int first = indices[0]; + for (int j = 1; j < count; ++j) { + if (!ggml_ext_graph_delay(graph, first, j, indices[j])) + break; + first = indices[j] - j; + } + } + } + static bool may_alias(const ggml_tensor* output, const ggml_tensor* input) { + 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; + } + // Unallocated inputs get allocator storage, separate from a preallocated output. + if (!input->data) + return false; + if (!input->buffer) + return true; + if (output->buffer != input->buffer) + return false; + const auto dst = reinterpret_cast(output->data); + const auto src = reinterpret_cast(input->data) + offset; + return dst < src + size && src < dst + ggml_nbytes(output); + } + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) { + const auto* output = nodes[count - 1]; + // Fresh outputs cannot alias external inputs except the residual's elementwise reuse. + const bool fresh = !output->data && !output->view_src && + (output->op == GGML_OP_CUSTOM || output->op == GGML_OP_GLU || output->op == GGML_OP_CONCAT || + (output->op == GGML_OP_ADD && count > 1)); + *bytes = fresh ? 0 : ggml_nbytes(output); + if (output->data && output->buffer && !output->view_src) { + bool alias = false; + for (int i = 0; i < count; ++i) + for (const auto* src : nodes[i]->src) + if (src && std::find(nodes, nodes + count, src) == nodes + count) + alias |= may_alias(output, src); + if (!alias) + *bytes = 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 count, void*, void*) { + Device d; + if (!device_info(ggml_backend_get_device(e->backend), d)) + return GGML_STATUS_FAILED; + return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count); + } +#undef SD_VK_PIPELINE +} + +void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) { + using namespace sd_vulkan_i8; + Device info; + if (!device_info(device, info)) + return; + ggml_ext_kernel_desc d = {}; + d.abi_version = GGML_EXT_ABI_VERSION; + d.priority = 100; + d.prepare = prepare; + d.destroy = destroy; + d.execute = execute; + d.name = "sd.quantize_i8_convrot.vulkan.v1"; + d.op = GGML_OP_CUSTOM; + d.custom_op = sd_ggml_convrot_op(); + d.supports = quant_supports; + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + d.name = "sd.linear_i8.vulkan.v1"; + d.op = GGML_OP_CUSTOM; + d.custom_op = sd_ggml_i8_op(); + d.supports = mat_supports; + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + d.priority = 200; + d.supports = nullptr; + d.match = match; + d.optimize = optimize; + d.name = "sd.linear_i8.fused.vulkan.v1"; + d.op = GGML_OP_CUSTOM; + d.custom_op = sd_ggml_i8_op(); + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + d.optimize = nullptr; + d.name = "sd.modulate_convrot.vulkan.v1"; + d.custom_op = sd_ggml_modulate_op(); + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + d.custom_op = nullptr; + const struct { + ggml_op op; + const char* name; + } fusions[] = { + {GGML_OP_NORM, "sd.norm_convrot.vulkan.v1"}, + {GGML_OP_RMS_NORM, "sd.rms_modulate_convrot.vulkan.v1"}, + {GGML_OP_GLU, "sd.swiglu_convrot.vulkan.v1"}, + {GGML_OP_UNARY, "sd.sigmoid_convrot.vulkan.v1"}}; + for (const auto& fusion : fusions) { + d.op = fusion.op; + d.name = fusion.name; + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + } +} diff --git a/src/core/vulkan/modulate.cpp b/src/core/vulkan/modulate.cpp new file mode 100644 index 0000000000..4561447022 --- /dev/null +++ b/src/core/vulkan/modulate.cpp @@ -0,0 +1,134 @@ +#include "modulate.h" +#include +#include +#include "core/ggml_extensions.h" +#include "ggml-ext-vulkan.h" + +namespace sd_vulkan_modulate { + struct Parameters { + uint32_t ne[4], scale_ne[4], shift_ne[4]; + uint32_t x_stride[4], scale_stride[4], shift_stride[4]; + uint32_t offsets[4]; + uint32_t count, has_shift; + }; + static_assert(sizeof(Parameters) == 120, "Modulate push constant layout mismatch"); + + static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) { + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1")); + return query ? query() : nullptr; + } + + static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info) || !sd_ggml_modulate_describe(node) || + info.max_group_size[0] < 256 || info.max_group_invocations < 256 || info.max_push_constants < sizeof(Parameters) || ggml_nelements(node) > UINT32_MAX - 255) { + return false; + } + const ggml_tensor* tensors[] = {node, node->src[0], node->src[1], node->src[2]}; + for (const auto* t : tensors) { + if (!t) { + continue; + } + size_t offset = 0; + const auto* owner = t; + while (!owner->data && owner->view_src) { + offset += owner->view_offs; + owner = owner->view_src; + } + if ((reinterpret_cast(owner->data) + offset) % 4 || + ggml_nbytes(t) > info.max_buffer_range - (info.buffer_alignment - 1)) { + return false; + } + for (int d = 0; d < 4; ++d) { + if (t->nb[d] / 4 > UINT32_MAX || t->ne[d] > UINT32_MAX) { + return false; + } + } + } + return (uint64_t(ggml_nelements(node)) + 255) / 256 <= uint64_t(info.max_group_count[0]) * info.max_group_count[1]; + } + + 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*) { + auto device = ggml_backend_get_device(e->backend); + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info)) { + return GGML_STATUS_FAILED; + } + const auto* dst = nodes[0]; + const auto* x = dst->src[0]; + const auto* scale = dst->src[1]; + const auto* shift = dst->src[2]; + ggml_ext_vulkan_binding bindings[] = {{x}, {scale}, {shift ? shift : scale}, {dst}}; + ggml_ext_vulkan_buffer_info buffers[4]; + for (int i = 0; i < 4; ++i) { + if (!service->buffer_info(e->backend, &bindings[i], &buffers[i])) { + return GGML_STATUS_FAILED; + } + } + bool alias = false; + for (int i = 0; i < 3; ++i) { + alias |= buffers[i].storage == buffers[3].storage && buffers[i].offset < buffers[3].offset + buffers[3].size && + buffers[3].offset < buffers[i].offset + buffers[i].size; + } + const auto destination = bindings[3]; + if (alias) { + bindings[3] = {nullptr, e->workspace, 0, ggml_nbytes(dst)}; + if (!service->buffer_info(e->backend, &bindings[3], &buffers[3])) { + return GGML_STATUS_FAILED; + } + } + Parameters p = {}; + p.count = uint32_t(ggml_nelements(dst)); + p.has_shift = shift != nullptr; + for (int d = 0; d < 4; ++d) { + p.ne[d] = uint32_t(x->ne[d]); + p.scale_ne[d] = uint32_t(scale->ne[d]); + p.shift_ne[d] = shift ? uint32_t(shift->ne[d]) : 1; + p.x_stride[d] = uint32_t(x->nb[d] / 4); + p.scale_stride[d] = uint32_t(scale->nb[d] / 4); + p.shift_stride[d] = shift ? uint32_t(shift->nb[d] / 4) : 0; + p.offsets[d] = buffers[d].shader_offset / 4; + } + const uint32_t count = (p.count + 255) / 256; + const uint32_t width = std::min(count, info.max_group_count[0]); + const uint32_t groups[] = {width, (count + width - 1) / width, 1}; + ggml_ext_vulkan_pipeline pipeline = {}; + pipeline.name = "sd.modulate.vulkan"; + pipeline.spirv = sd_vulkan_shaders::modulate; + pipeline.spirv_size = sizeof(sd_vulkan_shaders::modulate); + pipeline.binding_count = 4; + pipeline.push_constant_size = sizeof(Parameters); + auto status = service->dispatch(e->backend, &pipeline, bindings, &p, groups); + if (status == GGML_STATUS_SUCCESS && alias) { + status = service->copy(e->backend, &bindings[3], &destination); + } + return status; + } +} + +void sd_ggml_register_vulkan_modulate(ggml_backend_dev_t device) { + if (!sd_vulkan_modulate::api(device)) { + return; + } + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.modulate.vulkan.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_modulate_op(); + desc.priority = 100; + desc.supports = sd_vulkan_modulate::supports; + desc.prepare = sd_vulkan_modulate::prepare; + desc.destroy = sd_vulkan_modulate::destroy; + desc.execute = sd_vulkan_modulate::execute; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/vulkan/rms_affine.cpp b/src/core/vulkan/rms_affine.cpp new file mode 100644 index 0000000000..2eee7c7e37 --- /dev/null +++ b/src/core/vulkan/rms_affine.cpp @@ -0,0 +1,192 @@ +#include +#include +#include +#include "core/ggml_extensions.h" +#include "ggml-ext-vulkan.h" +#include "rms_modulate.h" + +namespace sd_vulkan_rms_affine { + struct Parameters { + uint32_t ne[4], stride[4], offsets[5]; + float eps; + uint32_t ne10; + }; + static_assert(sizeof(Parameters) == 60, "RMS affine push constant layout"); + + static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) { + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1")); + return query ? query() : nullptr; + } + static bool bounded(const ggml_tensor* t, const ggml_ext_vulkan_device_info& info) { + if (!t || t->type != GGML_TYPE_F32 || t->nb[0] != 4 || ggml_nbytes(t) > info.max_buffer_range - (info.buffer_alignment - 1)) + return false; + size_t offset = 0; + const auto* owner = t; + while (!owner->data && owner->view_src) { + offset += owner->view_offs; + owner = owner->view_src; + } + if ((reinterpret_cast(owner->data) + offset) % 4) + return false; + for (int d = 0; d < 4; ++d) + if (t->ne[d] > UINT32_MAX || t->nb[d] % 4 || t->nb[d] / 4 > UINT32_MAX) + return false; + return true; + } + static const ggml_tensor* other(const ggml_tensor* node, const ggml_tensor* input) { + return node->src[0] == input ? node->src[1] : node->src[1] == input ? node->src[0] + : nullptr; + } + struct Operation { + const ggml_tensor* x; + const ggml_tensor* weight; + const ggml_tensor* scale; + const ggml_tensor* shift; + const ggml_tensor* output; + }; + static Operation describe(ggml_tensor* const* nodes, int count) { + Operation op = {nodes[0]->src[0], other(nodes[1], nodes[0]), nullptr, nullptr, nodes[count - 1]}; + if (count == 3) { + op.scale = nodes[2]->src[1]; + op.shift = nodes[2]->src[2]; + } else if (count == 5) { + op.scale = other(nodes[2], nodes[1]); + op.shift = other(nodes[4], nodes[3]); + } + return op; + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + auto* g = const_cast(graph); + if (!service || !service->device_info(device, &info) || info.fusion_disabled || first + 1 >= ggml_graph_n_nodes(g)) + return 0; + ggml_tensor* nodes[5] = {}; + const int available = std::min(5, ggml_graph_n_nodes(g) - first); + for (int i = 0; i < available; ++i) + nodes[i] = ggml_graph_node(g, first + i); + if (nodes[0]->op != GGML_OP_RMS_NORM || nodes[1]->op != GGML_OP_MUL) + return 0; + const auto* x = nodes[0]->src[0]; + if (!bounded(x, info) || x->ne[0] != 3072 || available < 3 || !ggml_is_contiguous(x) || + x->ne[1] <= 1 || x->ne[1] > info.max_group_count[0] || x->ne[2] != 1 || x->ne[3] != 1 || + info.max_group_size[0] < 512 || info.max_group_invocations < 512) + return 0; + int count; + if (sd_ggml_modulate_describe(nodes[2]) && nodes[2]->src[0] == nodes[1]) { + count = 3; + } else if (available == 5 && nodes[2]->op == GGML_OP_MUL && nodes[3]->op == GGML_OP_ADD && + other(nodes[3], nodes[1]) == nodes[2] && nodes[4]->op == GGML_OP_ADD) { + count = 5; + } else { + return 0; + } + const auto op = describe(nodes, count); + if (!op.weight || !op.scale || (count == 5 && !op.shift)) + return 0; + for (const auto* parameter : {op.weight, op.scale, op.shift}) { + if (!parameter) + continue; + if (!bounded(parameter, info) || !ggml_is_contiguous(parameter) || parameter->ne[0] != x->ne[0] || ggml_nelements(parameter) != x->ne[0]) + return 0; + for (int i = 0; i < count; ++i) + if (parameter == nodes[i]) + return 0; + } + for (int i = 0; i < count; ++i) + if (!bounded(nodes[i], info) || !ggml_are_same_shape(nodes[i], x) || !ggml_is_contiguous(nodes[i])) + return 0; + return count; + } + static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info) || info.fusion_disabled) + return; + for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) { + const auto* node = ggml_graph_node(graph, i); + if (!sd_ggml_modulate_describe(node) && node->op != GGML_OP_MUL) + continue; + const auto* weighted = node->src[0]; + if (!weighted || weighted->op != GGML_OP_MUL || !weighted->src[0] || + weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 3072) + continue; + for (int j = i - 2; j >= std::max(0, i - 32); --j) { + if (ggml_graph_node(graph, j) == weighted->src[0] && ggml_graph_node(graph, j + 1) == weighted) { + if (j + 2 < i) + ggml_ext_graph_delay(graph, j, 2, i); + break; + } + } + } + } + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) { + *bytes = ggml_nbytes(nodes[count - 1]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void*, void*) {} + static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) { + auto device = ggml_backend_get_device(e->backend); + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info)) + return GGML_STATUS_FAILED; + const auto op = describe(nodes, count); + ggml_ext_vulkan_binding bindings[] = {{op.x}, {op.weight}, {op.output}, {op.scale}, {op.shift ? op.shift : op.scale}}; + ggml_ext_vulkan_buffer_info buffers[5]; + const int n = 5; + for (int i = 0; i < n; ++i) + if (!service->buffer_info(e->backend, bindings + i, buffers + i)) + return GGML_STATUS_FAILED; + bool alias = false; + for (int i = 0; i < n; ++i) + if (i != 2) + alias |= buffers[i].storage == buffers[2].storage && buffers[i].offset < buffers[2].offset + buffers[2].size && + buffers[2].offset < buffers[i].offset + buffers[i].size; + const auto destination = bindings[2]; + if (alias) { + bindings[2] = {nullptr, e->workspace, 0, ggml_nbytes(op.output)}; + if (!service->buffer_info(e->backend, bindings + 2, buffers + 2)) + return GGML_STATUS_FAILED; + } + Parameters p = {}; + p.ne10 = uint32_t(op.weight->ne[0]); + for (int i = 0; i < 4; ++i) { + p.ne[i] = uint32_t(op.x->ne[i]); + p.stride[i] = uint32_t(op.x->nb[i] / 4); + } + for (int i = 0; i < n; ++i) + p.offsets[i] = buffers[i].shader_offset / 4; + std::memcpy(&p.eps, nodes[0]->op_params, sizeof(p.eps)); + const uint32_t has_shift = op.shift != nullptr; + ggml_ext_vulkan_pipeline pipeline = {}; + pipeline.name = "sd.rms_modulate.vulkan"; + pipeline.spirv = sd_vulkan_shaders::rms_modulate; + pipeline.spirv_size = sizeof(sd_vulkan_shaders::rms_modulate); + pipeline.binding_count = n; + pipeline.push_constant_size = sizeof(p); + pipeline.constants = &has_shift; + pipeline.constant_count = 1; + uint32_t groups[] = {p.ne[1], std::min(p.ne[2], info.max_group_count[1]), std::min(p.ne[3], info.max_group_count[2])}; + auto status = service->dispatch(e->backend, &pipeline, bindings, &p, groups); + if (status == GGML_STATUS_SUCCESS && alias) + status = service->copy(e->backend, bindings + 2, &destination); + return status; + } +} +void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device) { + if (!sd_vulkan_rms_affine::api(device)) + return; + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rms_affine.vulkan.v1"; + desc.op = GGML_OP_RMS_NORM; + desc.priority = 150; + desc.match = sd_vulkan_rms_affine::match; + desc.optimize = sd_vulkan_rms_affine::optimize; + desc.prepare = sd_vulkan_rms_affine::prepare; + desc.execute = sd_vulkan_rms_affine::execute; + desc.destroy = sd_vulkan_rms_affine::destroy; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/vulkan/rope.cpp b/src/core/vulkan/rope.cpp new file mode 100644 index 0000000000..194603e13d --- /dev/null +++ b/src/core/vulkan/rope.cpp @@ -0,0 +1,234 @@ +#include "core/ggml_extensions.h" +#include "ggml-ext-vulkan.h" +#include "rms_rope_apply.h" +#include "rms_rope_cast.h" +#include "rope_apply.h" + +#include +#include +#include + +namespace sd_vulkan_rope { + struct Parameters { + uint32_t half_dim, rot_half, heads, tokens, pairs; + uint32_t x_stride[4], pe_stride[4]; + uint32_t x_offset, pe_offset, dst_offset, interleaved; + }; + struct FusedParameters { + Parameters rope; + float eps; + uint32_t w_offset; + float output_scale; + }; + static_assert(sizeof(Parameters) == 68 && sizeof(FusedParameters) == 80, "RoPE push constant layout mismatch"); + static const ggml_ext_vulkan_api* api(ggml_backend_dev_t device) { + auto query = reinterpret_cast(ggml_backend_reg_get_proc_address( + ggml_backend_dev_backend_reg(device), "ggml_ext_vulkan_v1")); + return query ? query() : nullptr; + } + static float param(const ggml_tensor* node, int i) { + float value; + std::memcpy(&value, node->op_params + i, sizeof(value)); + return value; + } + static bool is_rope(const ggml_tensor* node) { + return node && ggml_ext_op_get(node) == sd_ggml_rope_op(); + } + 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(tensor->data) + offset) % alignment == 0; + } + static bool bounded(const ggml_tensor* node, const ggml_ext_vulkan_device_info& info) { + if (!node || !aligned(node, 4) || ggml_nbytes(node) > info.max_buffer_range - (info.buffer_alignment - 1)) + return false; + for (int i = 0; i < 4; ++i) + if (node->nb[i] % 4 || node->nb[i] / 4 > UINT32_MAX || node->ne[i] > UINT32_MAX) + return false; + return true; + } + static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info) || info.max_group_size[0] < 256 || info.max_group_invocations < 256) + return false; + const auto* x = node->src[0]; + const auto* pe = node->src[1]; + return x && pe && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 && node->type == GGML_TYPE_F32 && + x->ne[0] % 2 == 0 && 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] && node->ne[0] == x->ne[0] && node->ne[1] == x->ne[2] && + node->ne[2] == x->ne[1] * x->ne[3] && node->ne[3] == 1 && ggml_is_contiguous(node) && + bounded(x, info) && bounded(pe, info) && bounded(node, info) && ggml_nelements(node) / 2 <= UINT32_MAX && + (uint64_t(ggml_nelements(node) / 2) + 255) / 256 <= uint64_t(info.max_group_count[0]) * info.max_group_count[1]; + } + static bool fused_supports(ggml_backend_dev_t device, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) { + ggml_ext_vulkan_device_info info; + if (!api(device)->device_info(device, &info) || info.fusion_disabled || !info.subgroup_shuffle || !info.full_subgroups || + info.subgroup_size != 32 || !is_rope(rope) || !supports(device, rope, nullptr) || + norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || mul->src[0] != norm || rope->src[0] != mul) + return false; + const auto* x = norm->src[0]; + const auto* weight = mul->src[1]; + return x && weight && weight != norm && x->type == GGML_TYPE_F32 && x->ne[0] == 128 && x->nb[0] == 4 && + norm->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_are_same_shape(x, norm) && ggml_are_same_shape(norm, mul) && + weight->type == GGML_TYPE_F32 && weight->ne[0] == 128 && ggml_nelements(weight) == 128 && + ggml_is_contiguous(weight) && bounded(weight, info) && bounded(x, info) && (x->ne[1] + 3) / 4 <= info.max_group_count[0]; + } + static bool cast_supports(ggml_backend_dev_t device, const ggml_tensor* rope, const ggml_tensor* scale, const ggml_tensor* copy) { + ggml_ext_vulkan_device_info info; + api(device)->device_info(device, &info); + // Workspace writeback uses Vulkan's four-byte copy granularity. + return info.float16 && (!scale || (scale->op == GGML_OP_SCALE && scale->src[0] == rope && scale->type == GGML_TYPE_F32 && param(scale, 1) == 0.f)) && + copy->op == GGML_OP_CPY && copy->src[0] == (scale ? scale : rope) && copy->type == GGML_TYPE_F16 && + ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(rope, copy); + } + static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (first + 2 >= ggml_graph_n_nodes(g) || !fused_supports(device, 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) && cast_supports(device, ggml_graph_node(g, first + 2), ggml_graph_node(g, first + 3), ggml_graph_node(g, first + 4))) + return 5; + return first + 3 < ggml_graph_n_nodes(g) && cast_supports(device, ggml_graph_node(g, first + 2), nullptr, ggml_graph_node(g, first + 3)) ? 4 : 3; + } + static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) { + for (int last = 2; last < ggml_graph_n_nodes(graph); ++last) { + const auto* output = ggml_graph_node(graph, last); + const auto* source = output->op == GGML_OP_CPY ? output->src[0] : nullptr; + const auto* scale = source && source->op == GGML_OP_SCALE ? source : nullptr; + const auto* rope = source ? (scale ? scale->src[0] : source) : output; + if (!is_rope(rope)) + continue; + const auto* mul = rope->src[0]; + const auto* norm = mul->src[0]; + if (!norm || !fused_supports(device, norm, mul, rope)) + continue; + const int count = output == rope ? 3 : cast_supports(device, rope, scale, output) ? (scale ? 5 : 4) + : 0; + if (!count) + continue; + const ggml_tensor* nodes[] = {norm, mul, rope, scale ? scale : output, output}; + int indices[5], next = count - 1; + for (int i = last; i >= std::max(0, last - 64) && next >= 0; --i) + if (ggml_graph_node(graph, i) == nodes[next]) + indices[next--] = i; + if (next >= 0) + continue; + int first = indices[0]; + for (int i = 1; i < count; ++i) { + if (!ggml_ext_graph_delay(graph, first, i, indices[i])) + break; + first = indices[i] - i; + } + } + } + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) { + const auto* output = nodes[count - 1]; + // RoPE outputs and owning casts cannot reuse allocator input storage. + *bytes = output->data || output->view_src ? ggml_nbytes(output) : 0; + return GGML_STATUS_SUCCESS; + } + static void destroy(void*, void*) {} + static bool overlaps(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) { + return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size; + } + static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) { + auto device = ggml_backend_get_device(execution->backend); + const auto* service = api(device); + ggml_ext_vulkan_device_info info; + if (!service || !service->device_info(device, &info)) + return GGML_STATUS_FAILED; + const auto* rope = nodes[count == 1 ? 0 : 2]; + const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0]; + const auto* pe = rope->src[1]; + const auto* output = nodes[count - 1]; + ggml_ext_vulkan_binding bindings[] = {{x}, {pe}, {output}, {count == 1 ? nullptr : nodes[1]->src[1]}}; + ggml_ext_vulkan_buffer_info buffers[4]; + const int n_buffers = count == 1 ? 3 : 4; + for (int i = 0; i < n_buffers; ++i) + if (!service->buffer_info(execution->backend, bindings + i, buffers + i)) + return GGML_STATUS_FAILED; + bool alias = overlaps(buffers[2], buffers[0]) || overlaps(buffers[2], buffers[1]) || (count != 1 && overlaps(buffers[2], buffers[3])); + const auto destination = bindings[2]; + if (alias) { + bindings[2] = {nullptr, execution->workspace, 0, ggml_nbytes(output)}; + if (!service->buffer_info(execution->backend, bindings + 2, buffers + 2)) + return GGML_STATUS_FAILED; + } + FusedParameters params = {}; + auto& pc = params.rope; + pc.half_dim = uint32_t(x->ne[0] / 2); + pc.rot_half = uint32_t(pe->ne[2]); + pc.heads = uint32_t(x->ne[1]); + pc.tokens = uint32_t(x->ne[2]); + pc.pairs = uint32_t(ggml_nelements(x) / 2); + for (int i = 0; i < 4; ++i) { + pc.x_stride[i] = uint32_t(x->nb[i] / 4); + pc.pe_stride[i] = uint32_t(pe->nb[i] / 4); + } + pc.x_offset = buffers[0].shader_offset / 4; + pc.pe_offset = buffers[1].shader_offset / 4; + pc.dst_offset = buffers[2].shader_offset / uint32_t(ggml_type_size(output->type)); + std::memcpy(&pc.interleaved, ggml_ext_op_attributes(rope), 4); + uint32_t groups[3] = {1, 1, 1}; + const uint32_t apply_output_scale = count == 5; + ggml_ext_vulkan_pipeline pipeline = {}; + pipeline.name = "sd.rope_apply.vulkan"; + pipeline.spirv = sd_vulkan_shaders::rope_apply; + pipeline.spirv_size = sizeof(sd_vulkan_shaders::rope_apply); + pipeline.binding_count = n_buffers; + pipeline.push_constant_size = sizeof(Parameters); + if (count == 1) { + const uint32_t n_groups = uint32_t((uint64_t(pc.pairs) + 255) / 256); + groups[0] = std::min(n_groups, info.max_group_count[0]); + groups[1] = (n_groups + groups[0] - 1) / groups[0]; + } else { + pipeline.name = count > 3 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan"; + pipeline.spirv = count > 3 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply; + pipeline.spirv_size = count > 3 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply); + if (count > 3) { + pipeline.constants = &apply_output_scale; + pipeline.constant_count = 1; + } + pipeline.push_constant_size = sizeof(FusedParameters); + pipeline.full_subgroups = true; + pipeline.required_subgroup_size = info.subgroup_size_control ? 32 : 0; + params.eps = param(nodes[0], 0); + params.w_offset = buffers[3].shader_offset / 4; + params.output_scale = count == 5 ? param(nodes[3], 0) : 1.f; + groups[0] = (pc.heads + 3) / 4; + groups[1] = std::min(pc.tokens, info.max_group_count[1]); + groups[2] = std::min(uint32_t(x->ne[3]), info.max_group_count[2]); + } + auto status = service->dispatch(execution->backend, &pipeline, bindings, ¶ms, groups); + if (status == GGML_STATUS_SUCCESS && alias) + status = service->copy(execution->backend, bindings + 2, &destination); + return status; + } +} + +void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device) { + if (!sd_vulkan_rope::api(device)) + return; + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.rope_apply.vulkan.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_rope_op(); + desc.priority = 100; + desc.supports = sd_vulkan_rope::supports; + desc.prepare = sd_vulkan_rope::prepare; + desc.execute = sd_vulkan_rope::execute; + desc.destroy = sd_vulkan_rope::destroy; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + desc.name = "sd.rms_rope_apply.vulkan.v1"; + desc.op = GGML_OP_RMS_NORM; + desc.custom_op = nullptr; + desc.priority = 200; + desc.supports = nullptr; + desc.match = sd_vulkan_rope::match; + desc.optimize = sd_vulkan_rope::optimize; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); +} diff --git a/src/core/vulkan/shaders/modulate.comp b/src/core/vulkan/shaders/modulate.comp new file mode 100644 index 0000000000..81f910d07b --- /dev/null +++ b/src/core/vulkan/shaders/modulate.comp @@ -0,0 +1,48 @@ +#version 450 + +layout(local_size_x = 256) in; +layout(binding = 0) readonly buffer X { + float x[]; +}; +layout(binding = 1) readonly buffer S { + float scale[]; +}; +layout(binding = 2) readonly buffer H { + float shift[]; +}; +layout(binding = 3) writeonly buffer D { + float dst[]; +}; +layout(push_constant) uniform Parameters { + uint ne[4], scale_ne[4], shift_ne[4]; + uint x_stride[4], scale_stride[4], shift_stride[4]; + uint offsets[4]; + uint count, has_shift; +} +p; + +void main() { + uint group = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; + if (group >= (p.count + 255) / 256) { + return; + } + uint i = group * 256 + gl_LocalInvocationID.x; + if (i >= p.count) { + return; + } + uint index = i; + uint xi = p.offsets[0], si = p.offsets[1], hi = p.offsets[2]; + for (uint d = 0; d < 4; ++d) { + uint 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]; + hi += (coord % p.shift_ne[d]) * p.shift_stride[d]; + } + precise float product = x[xi] * scale[si]; + precise float value = x[xi] + product; + if (p.has_shift != 0) { + value += shift[hi]; + } + dst[p.offsets[3] + i] = value; +} diff --git a/src/core/vulkan/shaders/modulate_convrot.comp b/src/core/vulkan/shaders/modulate_convrot.comp new file mode 100644 index 0000000000..4e85b1d0b8 --- /dev/null +++ b/src/core/vulkan/shaders/modulate_convrot.comp @@ -0,0 +1,163 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +layout(constant_id = 0) const uint K = 4096; +layout(constant_id = 1) const bool AFFINE = false; +layout(constant_id = 2) const bool HAS_SHIFT = true; +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; +layout(binding = 0) readonly buffer X { + float data_x[]; +}; +layout(binding = 1) readonly buffer M { + float data_m[]; +}; +layout(binding = 2) readonly buffer H { + float data_h[]; +}; +layout(binding = 3) writeonly buffer D { + uint data_d[]; +}; +layout(binding = 4) writeonly buffer S { + float data_s[]; +}; +#ifdef RMS_AFFINE +layout(binding = 5) readonly buffer W { + float data_w[]; +}; +#endif +layout(push_constant) uniform parameter { + uint rows; + uint x_offset; + uint m_offset; + uint h_offset; + float eps; + uint k; + uint w_offset; +} +p; +shared float values[K]; +shared vec2 sums[512]; + +uint quantize(float value, float inv_scale) { + return uint(clamp(int(roundEven(value * inv_scale)), -127, 127)) & 0xffu; +} + +#ifdef RMS_AFFINE +// Keep RMS arithmetic outside the scope of precise on the quantized output. +float rms_row(uint row, uint tid) { + float sum = 0.0; + [[unroll]] for (uint col = tid, iter = 0; iter < 12; col += 512, ++iter) { + float x = 0.0; + if (col < p.k) { + x = data_x[p.x_offset + row * K + col]; + values[col] = x; + } + sum += x * x; + } + sums[tid].x = sum; + barrier(); + [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) { + if (tid < stride) { + sum += sums[tid + stride].x; + sums[tid].x = sum; + } + barrier(); + } + return inversesqrt(sums[0].x / p.k + p.eps); +} + +float rms_weight(float scale, float value, float weight) { + return scale * value * weight; +} +#endif + +void main() { + const uint tid = gl_LocalInvocationID.x; + for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) { +#ifdef RMS_AFFINE + float inv_std = rms_row(row, tid); + float mean = 0.0; +#else + vec2 sum = vec2(0.0); + [[unroll]] for (uint col = tid; col < K; col += 512) { + const float x = data_x[p.x_offset + row * K + col]; + values[col] = x; + if (!AFFINE) { + sum.x += x; + sum.y = fma(x, x, sum.y); + } + } + float mean = 0.0; + float inv_std = 1.0; + if (!AFFINE) { + sums[tid] = sum; + barrier(); + [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) { + if (tid < stride) + sums[tid] += sums[tid + stride]; + barrier(); + } + mean = sums[0].x / p.k; + // Match NORM's contraction even when precise propagates from the quantized output. + const float variance = fma(-mean, mean, sums[0].y / p.k); + inv_std = inversesqrt(variance + p.eps); + } +#endif + [[unroll]] for (uint col = tid; col < K; col += 512) { + precise float value; +#ifdef RMS_AFFINE + values[col] = rms_weight(inv_std, values[col], data_w[p.w_offset + col]); +#endif + if (AFFINE) { + precise float modulated = values[col] * data_m[p.m_offset + col]; + value = values[col] + modulated; + if (HAS_SHIFT) { + value += data_h[p.h_offset + col]; + } + } else { + value = (values[col] - mean) * inv_std; + value *= data_m[p.m_offset + col]; + } + values[col] = value * (1.0 / 16.0); + } + barrier(); + [[unroll]] for (uint stride = 1; stride < 256; stride *= 4) { + for (uint i = tid; i < K / 4; i += 512) { + const uint group = i / 64; + const uint lane = i % 64; + const uint base = group * 256 + (lane / stride) * 4 * stride + lane % stride; + const float a = values[base]; + const float b = values[base + stride]; + const float c = values[base + 2 * stride]; + const float d = values[base + 3 * stride]; + values[base] = a + b + c - d; + values[base + stride] = a + b - c + d; + values[base + 2 * stride] = a - b + c + d; + values[base + 3 * stride] = -a + b + c + d; + } + barrier(); + } + float maximum = 0.0; + [[unroll]] for (uint col = tid; col < K; col += 512) { + maximum = max(maximum, abs(values[col])); + } + sums[tid].x = maximum; + barrier(); + [[unroll]] for (uint stride = 256; stride > 0; stride >>= 1) { + if (tid < stride) { + sums[tid].x = max(sums[tid].x, sums[tid + stride].x); + } + barrier(); + } + const float scale = sums[0].x / 127.0; + const float inv_scale = scale == 0.0 ? 0.0 : 1.0 / scale; + if (tid == 0) { + data_s[row] = scale; + } + for (uint i = tid; i < K / 4; i += 512) { + data_d[row * (K / 4) + i] = quantize(values[i * 4], inv_scale) | (quantize(values[i * 4 + 1], inv_scale) << 8) | (quantize(values[i * 4 + 2], inv_scale) << 16) | (quantize(values[i * 4 + 3], inv_scale) << 24); + } + barrier(); + } +} diff --git a/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp b/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp new file mode 100644 index 0000000000..7c0292e3ee --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp @@ -0,0 +1,74 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_integer_dot_product : require + +layout(local_size_x = 16, local_size_y = 16, local_size_z = 1) in; + +layout(push_constant) uniform parameter { + uint n; + uint rows; + uint k; + uint has_bias; + uint row_offset; +} +p; + +layout(binding = 0, std430) readonly buffer W { + int data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + int data_a[]; +}; +layout(binding = 2, std430) readonly buffer WS { + float data_ws[]; +}; +layout(binding = 3, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 4, std430) readonly buffer B { + float data_b[]; +}; +layout(binding = 5, std430) writeonly buffer D { + float data_d[]; +}; + +shared int weights[16 * 8]; +shared int activations[16 * 8]; + +void main() { + const uint output_index = gl_WorkGroupID.x * 16 + gl_LocalInvocationID.x; + const uint row = p.row_offset + gl_WorkGroupID.y * 16 + gl_LocalInvocationID.y; + const uint lx = gl_LocalInvocationID.x; + const uint ly = gl_LocalInvocationID.y; + const uint k_packed = p.k / 4; + int sum = 0; + + for (uint k_base = 0; k_base < k_packed; k_base += 8) { + if (ly < 8) { + weights[lx * 8 + ly] = output_index < p.n + ? data_w[output_index * k_packed + k_base + ly] + : 0; + } + if (lx < 8) { + activations[ly * 8 + lx] = row < p.rows + ? data_a[row * k_packed + k_base + lx] + : 0; + } + barrier(); + + [[unroll]] + for (uint i = 0; i < 8; ++i) { + sum += dotPacked4x8EXT(weights[lx * 8 + i], activations[ly * 8 + i]); + } + barrier(); + } + + if (output_index < p.n && row < p.rows) { + float value = float(sum) * data_as[row] * data_ws[output_index]; + if (p.has_bias != 0) { + value += data_b[output_index]; + } + data_d[row * p.n + output_index] = value; + } +} diff --git a/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp new file mode 100644 index 0000000000..573e257d16 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp @@ -0,0 +1,106 @@ +#version 460 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint BLOCK_SIZE = 32; +layout(constant_id = 1) const uint CM_M = 16; +layout(constant_id = 2) const uint CM_N = 16; +layout(constant_id = 3) const uint CM_K = 16; +const uint TILE_M = 4; +const uint TILE_N = 4; + +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform parameter { + uint n; + uint rows; + uint k; + uint has_bias; + uint row_offset; +} +p; + +layout(binding = 0, std430) readonly buffer W { + int8_t data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + int8_t data_a[]; +}; +layout(binding = 2, std430) readonly buffer WS { + float data_ws[]; +}; +layout(binding = 3, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 4, std430) readonly buffer B { + float data_b[]; +}; +layout(binding = 5, std430) writeonly buffer D { + float data_d[]; +}; + +shared int32_t accumulators[1024]; + +void main() { + const uint output_base = gl_WorkGroupID.x * CM_N * TILE_N; + const uint row_base = gl_WorkGroupID.y * CM_M * TILE_M; + coopmat sums[TILE_M * TILE_N]; + [[unroll]] + for (uint i = 0; i < TILE_M * TILE_N; ++i) { + sums[i] = coopmat(0); + } + + for (uint k_base = 0; k_base < p.k; k_base += CM_K) { + coopmat activations[TILE_M]; + coopmat weights[TILE_N]; + [[unroll]] + for (uint tile_m = 0; tile_m < TILE_M; ++tile_m) { + coopMatLoad(activations[tile_m], data_a, + (row_base + tile_m * CM_M) * p.k + k_base, + p.k, gl_CooperativeMatrixLayoutRowMajor); + } + [[unroll]] + for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) { + coopMatLoad(weights[tile_n], data_w, + (output_base + tile_n * CM_N) * p.k + k_base, + p.k, gl_CooperativeMatrixLayoutColumnMajor); + } + [[unroll]] + for (uint tile_m = 0; tile_m < TILE_M; ++tile_m) { + [[unroll]] + for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) { + const uint index = tile_m * TILE_N + tile_n; + sums[index] = coopMatMulAdd(activations[tile_m], weights[tile_n], sums[index]); + } + } + } + + [[unroll]] + for (uint tile_m = 0; tile_m < TILE_M; ++tile_m) { + [[unroll]] + for (uint tile_n = 0; tile_n < TILE_N; ++tile_n) { + coopMatStore(sums[tile_m * TILE_N + tile_n], accumulators, + 0, CM_N, gl_CooperativeMatrixLayoutRowMajor); + controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, + gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); + + for (uint index = gl_SubgroupInvocationID; index < CM_M * CM_N; index += gl_SubgroupSize) { + const uint row = row_base + tile_m * CM_M + index / CM_N; + const uint output_index = output_base + tile_n * CM_N + index % CM_N; + float value = float(accumulators[index]) * data_as[row] * data_ws[output_index]; + if (p.has_bias != 0) { + value += data_b[output_index]; + } + data_d[row * p.n + output_index] = value; + } + controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, + gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); + } + } +} diff --git a/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp new file mode 100644 index 0000000000..3cb4547c6b --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp @@ -0,0 +1,172 @@ +#version 460 + +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_NV_cooperative_matrix2 : require +#extension GL_KHR_memory_scope_semantics : require + +layout(constant_id = 0) const uint BN = 128; +layout(constant_id = 1) const uint ALIGNMENT = 0; +layout(constant_id = 2) const bool SWIGLU = false; +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; +layout(constant_id = 3) const uint BM = 128; +layout(constant_id = 4) const uint BK = 64; + +layout(push_constant) uniform parameter { + uint n; + uint rows; + uint k; + uint has_bias; + uint row_offset; + float output_scale; + uint group_cols; +} +p; + +layout(binding = 0, std430) readonly buffer W { + int8_t data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + int8_t data_a[]; +}; +layout(binding = 2, std430) readonly buffer WS { + float data_ws[]; +}; +layout(binding = 3, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 4, std430) readonly buffer B { + float data_b[]; +}; +// Store FP32 bit patterns without changing the integer accumulator fragment layout. +layout(binding = 5, std430) writeonly buffer D { + int32_t data_d[]; +}; +#ifdef SPLIT_SWIGLU +layout(binding = 6, std430) readonly buffer UW { + int8_t data_uw[]; +}; +layout(binding = 7, std430) readonly buffer US { + float data_us[]; +}; +#endif +#ifdef RESIDUAL +layout(binding = 6, std430) readonly buffer G { + float data_g[]; +}; +layout(binding = 7, std430) readonly buffer R { + float data_r[]; +}; +#endif +int32_t dequantize(const in uint row, const in uint col, const in int32_t acc, const in uint row_base, const in uint col_base) { + const uint r = row_base + col; + const uint c = col_base + row; + if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n)) { + return 0; + } + precise float value = float(acc) * data_as[r] * data_ws[c]; + if (p.has_bias != 0) { + value += data_b[c]; + } +#ifdef RESIDUAL + value *= p.output_scale; + value *= data_g[c]; + value += data_r[r * p.n + c]; +#endif + return floatBitsToInt(value); +} + +int32_t dequantize_swiglu(const in uint row, const in uint col, const in int32_t gate_acc, const in int32_t up_acc, const in uint row_base, const in uint col_base) { + const uint r = row_base + col; + const uint c = col_base + row; + if (ALIGNMENT != 1 && (r >= p.rows || c >= p.n / 2)) { + return 0; + } + precise float gate = float(gate_acc) * data_as[r] * data_ws[c]; +#ifdef SPLIT_SWIGLU + precise float up = float(up_acc) * data_as[r] * data_us[c]; +#else + precise float up = float(up_acc) * data_as[r] * data_ws[c + p.n / 2]; +#endif + return floatBitsToInt(gate / (1.0 + exp(-gate)) * up); +} + +void main() { + const uint k_size = p.k & ~255u; + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> a_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> w_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + a_layout = setTensorLayoutDimensionNV(a_layout, p.rows, k_size); + w_layout = setTensorLayoutDimensionNV(w_layout, p.n, k_size); + d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, SWIGLU ? p.n / 2 : p.n); + tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); + // Keep a weight panel in L2 while visiting the activation tiles. + const uint col_tiles = gl_NumWorkGroups.x; + const uint row_tiles = gl_NumWorkGroups.y; + const uint group_cols = clamp(p.group_cols, 1u, col_tiles); + const uint linear = gl_WorkGroupID.y * col_tiles + gl_WorkGroupID.x; + const uint group = linear / (group_cols * row_tiles); + const uint first_col = group * group_cols; + const uint cols = min(col_tiles - first_col, group_cols); + const uint within_group = linear - group * group_cols * row_tiles; + const uint row = p.row_offset + (within_group / cols) * BN; + const uint col = (first_col + within_group % cols) * BM; + coopmat sum = coopmat(0); + coopmat sum_up = coopmat(0); + if (ALIGNMENT != 0) { + tensorLayoutNV<2> a_full = createTensorLayoutNV(2); + tensorLayoutNV<2> w_full = createTensorLayoutNV(2); + a_full = setTensorLayoutDimensionNV(a_full, p.rows, k_size); + w_full = setTensorLayoutDimensionNV(w_full, p.n, k_size); + for (uint k = 0; k < p.k; k += BK) { + coopmat w; + coopmat a; + coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_full, col, BM, k, BK)); + if (ALIGNMENT == 1) { + coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(a_full, row, BN, k, BK), transpose); + } else { + coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(a_layout, row, BN, k, BK), transpose); + } + sum = coopMatMulAdd(w, a, sum); + if (SWIGLU) { +#ifdef SPLIT_SWIGLU + coopMatLoadTensorNV(w, data_uw, 0, sliceTensorLayoutNV(w_full, col, BM, k, BK)); +#else + coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_full, col + p.n / 2, BM, k, BK)); +#endif + sum_up = coopMatMulAdd(w, a, sum_up); + } + } + } else { + for (uint k = 0; k < p.k; k += BK) { + coopmat w; + coopmat a; + coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_layout, col, BM, k, BK)); + coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(a_layout, row, BN, k, BK), transpose); + sum = coopMatMulAdd(w, a, sum); + if (SWIGLU) { +#ifdef SPLIT_SWIGLU + coopMatLoadTensorNV(w, data_uw, 0, sliceTensorLayoutNV(w_layout, col, BM, k, BK)); +#else + coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(w_layout, col + p.n / 2, BM, k, BK)); +#endif + sum_up = coopMatMulAdd(w, a, sum_up); + } + } + } + if (SWIGLU) { + coopMatPerElementNV(sum, sum, dequantize_swiglu, sum_up, row, col); + } else { + coopMatPerElementNV(sum, sum, dequantize, row, col); + } + // The aligned dispatch contains only complete tiles and a 128-aligned output width. + if (ALIGNMENT == 1) { + tensorLayoutNV<2> d_full = createTensorLayoutNV(2); + d_full = setTensorLayoutDimensionNV(d_full, p.rows, (SWIGLU ? p.n / 2 : p.n) & ~127u); + coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_full, row, BN, col, BM), transpose); + } else { + coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row, BN, col, BM), transpose); + } +} diff --git a/src/core/vulkan/shaders/quantize_i8_convrot.comp b/src/core/vulkan/shaders/quantize_i8_convrot.comp new file mode 100644 index 0000000000..20edd6165e --- /dev/null +++ b/src/core/vulkan/shaders/quantize_i8_convrot.comp @@ -0,0 +1,105 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(push_constant) uniform parameter { + uint k; + uint rows; +} +p; + +layout(binding = 0, std430) readonly buffer A { + float data_a[]; +}; +layout(binding = 1, std430) writeonly buffer D { + uint data_d[]; +}; +layout(binding = 2, std430) writeonly buffer S { + float data_s[]; +}; + +shared float values[256]; +shared float maxima[64]; + +void transform_group(const uint row, const uint group) { + const uint tid = gl_LocalInvocationID.x; + const uint base_offset = row * p.k + group * 256; + + [[unroll]] + for (uint i = 0; i < 4; ++i) { + const uint index = tid + i * 64; + values[index] = data_a[base_offset + index] * (1.0 / 16.0); + } + barrier(); + + [[unroll]] + for (uint stride = 1; stride < 256; stride *= 4) { + const uint base = (tid / stride) * 4 * stride + tid % stride; + const uint i0 = base; + const uint i1 = i0 + stride; + const uint i2 = i1 + stride; + const uint 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; + barrier(); + } +} + +uint quantize_value(const float value, const float inv_scale) { + const int quantized = clamp(int(roundEven(value * inv_scale)), -127, 127); + return uint(quantized) & 0xffu; +} + +void quantize_row(const uint row) { + const uint tid = gl_LocalInvocationID.x; + const uint groups = p.k / 256; + float local_max = 0.0; + + for (uint group = 0; group < groups; ++group) { + transform_group(row, group); + [[unroll]] + for (uint i = 0; i < 4; ++i) { + local_max = max(local_max, abs(values[tid + i * 64])); + } + barrier(); + } + + maxima[tid] = local_max; + barrier(); + [[unroll]] + for (uint stride = 32; stride > 0; stride >>= 1) { + if (tid < stride) { + maxima[tid] = max(maxima[tid], maxima[tid + stride]); + } + barrier(); + } + + if (tid == 0) { + data_s[row] = maxima[0] / 127.0; + } + barrier(); + + const float row_scale = maxima[0] / 127.0; + const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; + for (uint group = 0; group < groups; ++group) { + transform_group(row, group); + const uint index = tid * 4; + const uint packed = quantize_value(values[index], inv_scale) | (quantize_value(values[index + 1], inv_scale) << 8) | (quantize_value(values[index + 2], inv_scale) << 16) | (quantize_value(values[index + 3], inv_scale) << 24); + data_d[(row * p.k + group * 256) / 4 + tid] = packed; + barrier(); + } +} + +void main() { + for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) { + quantize_row(row); + } +} diff --git a/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp b/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp new file mode 100644 index 0000000000..eb451f6b36 --- /dev/null +++ b/src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp @@ -0,0 +1,140 @@ +#version 450 + +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; +layout(push_constant) uniform parameter { + uint k; + uint rows; +} +p; +layout(binding = 0, std430) readonly buffer A { + float data_a[]; +}; +layout(binding = 1, std430) writeonly buffer D { + int8_t data_d[]; +}; +layout(binding = 2, std430) writeonly buffer S { + float data_s[]; +}; +#if defined(SWIGLU) || defined(SIGMOID_MUL) +layout(binding = 3, std430) readonly buffer B { + float data_b[]; +}; +#endif +shared float maxima[8]; +#ifdef SINGLE_PASS +#ifdef STATIC_K +const uint K = STATIC_K; +#else +layout(constant_id = 0) const uint K = 4096; +#endif +#endif + +vec4 rotate(float a, float b, float c, float d) { + precise vec4 r = vec4(a + b + c - d, a + b - c + d, + a - b + c + d, -a + b + c + d); + return r; +} + +void transform_group(uint row, uint group, out float v[8]) { + const uint lane = gl_SubgroupInvocationID; + const uint base = row * p.k + group * 256; + [[unroll]] for (uint j = 0; j < 8; ++j) { + float value = data_a[base + lane + j * 32]; +#ifdef SWIGLU + precise float activated = value / (1.0 + exp(-value)) * data_b[base + lane + j * 32]; + value = activated; +#endif +#ifdef SIGMOID_MUL + precise float gate = 1.0 / (1.0 + exp(-value)); + precise float activated = data_b[base + lane + j * 32] * gate; + value = activated; +#endif + v[j] = value * (1.0 / 16.0); + } + [[unroll]] for (uint stride = 1; stride <= 4; stride *= 4) { + const uint first = lane & ~(3u * stride); + [[unroll]] for (uint j = 0; j < 8; ++j) { + vec4 r = rotate(subgroupShuffle(v[j], first), subgroupShuffle(v[j], first + stride), + subgroupShuffle(v[j], first + 2 * stride), subgroupShuffle(v[j], first + 3 * stride)); + v[j] = r[(lane / stride) & 3u]; + } + } + [[unroll]] for (uint j = 0; j < 8; j += 2) { + vec4 r = rotate(subgroupShuffle(v[j], lane & 15u), subgroupShuffle(v[j], (lane & 15u) + 16), + subgroupShuffle(v[j + 1], lane & 15u), subgroupShuffle(v[j + 1], (lane & 15u) + 16)); + v[j] = r[lane / 16]; + v[j + 1] = r[2 + lane / 16]; + } + [[unroll]] for (uint j = 0; j < 2; ++j) { + vec4 r = rotate(v[j], v[j + 2], v[j + 4], v[j + 6]); + v[j] = r.x; + v[j + 2] = r.y; + v[j + 4] = r.z; + v[j + 6] = r.w; + } +} + +void main() { + for (uint row = gl_WorkGroupID.x; row < p.rows; row += gl_NumWorkGroups.x) { + float local_max = 0.0; +#ifdef SINGLE_PASS + float values[((K + 2047) / 2048) * 8]; + [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { + if (gl_SubgroupID + g * 8 >= K / 256) + continue; + float v[8]; + transform_group(row, gl_SubgroupID + g * 8, v); + [[unroll]] for (uint j = 0; j < 8; ++j) { + values[g * 8 + j] = v[j]; + local_max = max(local_max, abs(v[j])); + } + } +#else + for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { + float v[8]; + transform_group(row, group, v); + [[unroll]] for (uint j = 0; j < 8; ++j) + local_max = max(local_max, abs(v[j])); + } +#endif + local_max = subgroupMax(local_max); + if (gl_SubgroupInvocationID == 0) + maxima[gl_SubgroupID] = local_max; + barrier(); + if (gl_LocalInvocationID.x == 0) { + float row_max = 0.0; + for (uint j = 0; j < 8; ++j) + row_max = max(row_max, maxima[j]); + maxima[0] = row_max; + data_s[row] = row_max / 127.0; + } + barrier(); + const float row_scale = maxima[0] / 127.0; + const float inv_scale = row_scale == 0.0 ? 0.0 : 1.0 / row_scale; +#ifdef SINGLE_PASS + [[unroll]] for (uint g = 0; g < (K + 2047) / 2048; ++g) { + if (gl_SubgroupID + g * 8 >= K / 256) + continue; + [[unroll]] for (uint j = 0; j < 8; ++j) { + data_d[row * p.k + (gl_SubgroupID + g * 8) * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(values[g * 8 + j] * inv_scale)), -127, 127)); + } + } +#else + for (uint group = gl_SubgroupID; group < p.k / 256; group += 8) { + float v[8]; + transform_group(row, group, v); + [[unroll]] for (uint j = 0; j < 8; ++j) { + data_d[row * p.k + group * 256 + gl_SubgroupInvocationID + j * 32] = int8_t(clamp(int(roundEven(v[j] * inv_scale)), -127, 127)); + } + } +#endif + barrier(); + } +} diff --git a/src/core/vulkan/shaders/rms_modulate.comp b/src/core/vulkan/shaders/rms_modulate.comp new file mode 100644 index 0000000000..53d68bf2bc --- /dev/null +++ b/src/core/vulkan/shaders/rms_modulate.comp @@ -0,0 +1,136 @@ +#version 450 +#extension GL_EXT_control_flow_attributes : enable +#define BLOCK_SIZE 512 +#define FLOAT_TYPE float +#define D_TYPE float +layout(local_size_x = BLOCK_SIZE) in; +layout(binding = 0) readonly buffer X { + float data_a[]; +}; +layout(binding = 1) readonly buffer W { + float data_b[]; +}; +layout(binding = 2) writeonly buffer D { + float data_d[]; +}; +layout(binding = 3) readonly buffer S { + float data_c[]; +}; +layout(binding = 4) readonly buffer H { + float data_e[]; +}; +layout(constant_id = 0) const bool has_shift = true; +layout(constant_id = 1) const bool do_multiply = true; +layout(push_constant) uniform Parameters { + uint ne[4], stride[4], offsets[5]; + float eps; + uint ne10; +} +p; +shared float sumsh[BLOCK_SIZE]; +shared float modulation_input[BLOCK_SIZE]; +uint get_aoffset() { + return p.offsets[0]; +} +uint get_boffset() { + return p.offsets[1]; +} +uint get_doffset() { + return p.offsets[2]; +} +uint src1_idx(uint a, uint b, uint c, uint d) { + return 0; +} +uint fastmod(uint a, uint b) { + return a % b; +} +float modulate(float value, uint col) { + // Keep precise modulation from changing the native normalization's FMA contraction. + modulation_input[gl_LocalInvocationIndex] = value; + memoryBarrierShared(); + float weighted = modulation_input[gl_LocalInvocationIndex]; + precise float product = weighted * data_c[p.offsets[3] + col]; + precise float result = weighted + product; + if (has_shift) + result += data_e[p.offsets[4] + col]; + return result; +} +void rms_norm(uint num_iters) { + const uint ncols = p.ne[0]; + const uint nrows = gl_NumWorkGroups.x; + const uint nchannels = p.ne[2]; + const uint nsamples = p.ne[3]; + + const uint row = gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationID.x; + + const uint stride_row = p.stride[1]; + const uint stride_channel = p.stride[2]; + const uint stride_sample = p.stride[3]; + + // grid.y/z are clamped to the device workgroup limit, iterate over the excess channels/samples + for (uint samp = gl_WorkGroupID.z; samp < nsamples; samp += gl_NumWorkGroups.z) { + for (uint channel = gl_WorkGroupID.y; channel < nchannels; channel += gl_NumWorkGroups.y) { + barrier(); + + uint a_offset = samp * stride_sample + channel * stride_channel + row * stride_row + get_aoffset(); + uint b_offset = src1_idx(0, row, channel, samp) + get_boffset(); + uint d_offset = ((samp * nchannels + channel) * nrows + row) * ncols + get_doffset(); + FLOAT_TYPE sum = FLOAT_TYPE(0.0f); // partial sum for thread in warp + + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + FLOAT_TYPE xi = FLOAT_TYPE(0); + if (col < ncols) { + xi = FLOAT_TYPE(data_a[a_offset + col]); + } + sum += xi * xi; + } + + sumsh[tid] = sum; + // sum up partial sums and write back result + barrier(); + [[unroll]] for (int s = BLOCK_SIZE / 2; s > 0; s >>= 1) { + if (tid < s) { + sum += sumsh[tid + s]; + sumsh[tid] = sum; + } + barrier(); + } + sum = sumsh[0]; + + const FLOAT_TYPE mean = sum / FLOAT_TYPE(ncols); + const FLOAT_TYPE scale = inversesqrt(mean + FLOAT_TYPE(p.eps)); + + if (do_multiply) { + if (ncols > p.ne10) { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + fastmod(col, p.ne10)]); + data_d[d_offset + col] = modulate(value, col); + } + } else { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + FLOAT_TYPE value = scale * FLOAT_TYPE(data_a[a_offset + col]) * FLOAT_TYPE(data_b[b_offset + col]); + data_d[d_offset + col] = modulate(value, col); + } + } + } else { + [[unroll]] for (uint col = tid, idx = 0; idx < num_iters; col += BLOCK_SIZE, ++idx) { + if (col >= ncols) { + continue; + } + data_d[d_offset + col] = D_TYPE(scale * FLOAT_TYPE(data_a[a_offset + col])); + } + } + } + } +} + +void main() { + rms_norm(8); +} diff --git a/src/core/vulkan/shaders/rms_rope_apply.comp b/src/core/vulkan/shaders/rms_rope_apply.comp new file mode 100644 index 0000000000..ef2f581c73 --- /dev/null +++ b/src/core/vulkan/shaders/rms_rope_apply.comp @@ -0,0 +1,111 @@ +#version 450 + +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; +layout(binding = 0) readonly buffer X { + float data_x[]; +}; +layout(binding = 1) readonly buffer PE { + float data_pe[]; +}; +#ifdef CAST_F16 +layout(constant_id = 0) const bool apply_output_scale = false; +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +layout(binding = 2) writeonly buffer D { + float16_t data_d[]; +}; +#else +layout(binding = 2) writeonly buffer D { + float data_d[]; +}; +#endif +layout(binding = 3) readonly buffer W { + float data_w[]; +}; + +layout(push_constant) uniform parameter { + uint half_dim; + uint rot_half; + uint heads; + uint tokens; + uint pairs; + uint x_stride[4]; + uint pe_stride[4]; + uint x_offset; + uint pe_offset; + uint dst_offset; + uint interleaved; + float eps; + uint w_offset; + float output_scale; +} +pc; + +float shuffle_channel(vec4 values, uint channel) { + const uint lane = channel & 31u; + const vec4 shuffled = vec4(subgroupShuffle(values.x, lane), subgroupShuffle(values.y, lane), + subgroupShuffle(values.z, lane), subgroupShuffle(values.w, lane)); + return shuffled[channel / 32]; +} + +void main() { + const uint head = gl_WorkGroupID.x * 4 + gl_SubgroupID; + if (head >= pc.heads) { + return; + } + const uint lane = gl_SubgroupInvocationID; + const uint batches = pc.pairs / (64 * pc.heads * pc.tokens); + for (uint batch = gl_WorkGroupID.z; batch < batches; batch += gl_NumWorkGroups.z) { + for (uint token = gl_WorkGroupID.y; token < pc.tokens; token += gl_NumWorkGroups.y) { + const uint base = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; + vec4 x = vec4(data_x[base + lane], data_x[base + lane + 32], + data_x[base + lane + 64], data_x[base + lane + 96]); + precise vec4 squares = x * x; + precise float sum = (squares.x + squares.z) + (squares.y + squares.w); + [[unroll]] for (uint stride = 16; stride > 0; stride >>= 1) { + sum += subgroupShuffle(sum, lane ^ stride); + } + sum = subgroupShuffle(sum, 0); + const float scale = inversesqrt(sum / 128.0 + pc.eps); + precise vec4 weighted = scale * x; + [[unroll]] for (uint i = 0; i < 4; ++i) { + weighted[i] *= data_w[pc.w_offset + lane + i * 32]; + } + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint col = lane + i * 32; + precise float value = weighted[i]; + if (pc.interleaved != 0) { + const float a = subgroupShuffle(value, lane & ~1u); + const float b = subgroupShuffle(value, lane | 1u); + if (col / 2 < pc.rot_half) { + const uint matrix = pc.pe_offset + (col / 2) * pc.pe_stride[2] + token * pc.pe_stride[3] + (col & 1u) * pc.pe_stride[1]; + value = a * data_pe[matrix] + b * data_pe[matrix + pc.pe_stride[0]]; + } + } else { + const uint pair = col % pc.rot_half; + const float a = shuffle_channel(weighted, pair); + const float b = shuffle_channel(weighted, pair + pc.rot_half); + if (col < 2 * pc.rot_half) { + const uint matrix = pc.pe_offset + pair * pc.pe_stride[2] + token * pc.pe_stride[3]; + const float c = data_pe[matrix]; + const float s = data_pe[matrix + pc.pe_stride[1]]; + value = col < pc.rot_half ? a * c - b * s : a * s + b * c; + } + } +#ifdef CAST_F16 + if (apply_output_scale) { + value *= pc.output_scale; + value += 0.0; + } + data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = float16_t(value); +#else + data_d[pc.dst_offset + ((batch * pc.heads + head) * pc.tokens + token) * 128 + col] = value; +#endif + } + } + } +} diff --git a/src/core/vulkan/shaders/rope_apply.comp b/src/core/vulkan/shaders/rope_apply.comp new file mode 100644 index 0000000000..ccfc6d631a --- /dev/null +++ b/src/core/vulkan/shaders/rope_apply.comp @@ -0,0 +1,66 @@ +#version 450 + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer X { + float data_x[]; +}; +layout(binding = 1) readonly buffer PE { + float data_pe[]; +}; +layout(binding = 2) writeonly buffer D { + float data_d[]; +}; + +layout(push_constant) uniform parameter { + uint half_dim; + uint rot_half; + uint heads; + uint tokens; + uint pairs; + uint x_stride[4]; + uint pe_stride[4]; + uint x_offset; + uint pe_offset; + uint dst_offset; + uint interleaved; +} +pc; + +void main() { + const uint index = gl_GlobalInvocationID.x + gl_NumWorkGroups.x * gl_WorkGroupSize.x * gl_GlobalInvocationID.y; + if (index >= pc.pairs) { + return; + } + const uint j = index % pc.half_dim; + const uint row = index / pc.half_dim; + const uint token = row % pc.tokens; + const uint head = (row / pc.tokens) % pc.heads; + const uint batch = row / (pc.tokens * pc.heads); + const uint i0 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j : j; + const uint i1 = pc.interleaved != 0 || j >= pc.rot_half ? 2 * j + 1 : j + pc.rot_half; + const uint x_row = pc.x_offset + head * pc.x_stride[1] + token * pc.x_stride[2] + batch * pc.x_stride[3]; + const uint out_row = pc.dst_offset + row * pc.half_dim * 2; + const float a = data_x[x_row + i0 * pc.x_stride[0]]; + const float b = data_x[x_row + i1 * pc.x_stride[0]]; + if (j >= pc.rot_half) { + data_d[out_row + i0] = a; + data_d[out_row + i1] = b; + return; + } + const uint matrix = pc.pe_offset + j * pc.pe_stride[2] + token * pc.pe_stride[3]; + const float c = data_pe[matrix]; + const float s = data_pe[matrix + pc.pe_stride[1]]; + // Preserve the separate multiply/add rounding of the unfused graph. + precise float y0; + precise float y1; + if (pc.interleaved != 0) { + y0 = a * c + b * data_pe[matrix + pc.pe_stride[0]]; + y1 = a * s + b * data_pe[matrix + pc.pe_stride[0] + pc.pe_stride[1]]; + } else { + y0 = a * c - b * s; + y1 = a * s + b * c; + } + data_d[out_row + i0] = y0; + data_d[out_row + i1] = y1; +} diff --git a/src/detailer.cpp b/src/detailer.cpp index 72d01dc6f0..9c48ca4787 100644 --- a/src/detailer.cpp +++ b/src/detailer.cpp @@ -1,4 +1,6 @@ #include "detailer.h" +#include "generation_info.h" +#include "pipeline/request.h" #include #include @@ -291,7 +293,9 @@ static void filter_and_sort_detections(std::vector* detectio std::partial_sort(detections->begin(), detections->begin() + params.mask_k_largest, detections->end(), - [](const auto& a, const auto& b) { return detection_area(a) > detection_area(b); }); + [&](const auto& a, const auto& b) { + return params.mask_filter_by_confidence ? a.confidence > b.confidence : detection_area(a) > detection_area(b); + }); detections->resize(params.mask_k_largest); } @@ -545,7 +549,16 @@ static sd_image_t resize_crop_mask(const Mask& source, const CropRegion& crop, i return output; } -static sd_image_t copy_as_rgb(const sd_image_t& source) { +static sd_image_t copy_for_detailing(const sd_image_t& source) { + if (source.channel == 4) { + sd_image_t output = source; + size_t bytes = static_cast(source.width) * source.height * source.channel; + output.data = static_cast(malloc(bytes)); + if (output.data != nullptr) { + memcpy(output.data, source.data, bytes); + } + return output; + } CropRegion full{0, 0, static_cast(source.width), static_cast(source.height)}; return resize_crop_image(source, full, source.width, source.height); } @@ -566,7 +579,7 @@ static void composite_crop(sd_image_t* destination, continue; } float source_x = (x - crop.x1 + 0.5f) * scale_x - 0.5f; - size_t destination_index = (static_cast(y) * destination->width + x) * 3; + size_t destination_index = (static_cast(y) * destination->width + x) * destination->channel; for (int c = 0; c < 3; ++c) { float generated_value = bilinear_channel(generated, source_x, source_y, c) * 255.f; float original_value = destination->data[destination_index + c]; @@ -787,6 +800,10 @@ static bool parse_adetailer_params(const sd_adetailer_params_t& public_params, ok = parse_int(value, ¶ms->max_detections); } else if (key == "mask_k_largest") { ok = parse_int(value, ¶ms->mask_k_largest); + } else if (key == "mask_filter_method") { + const std::string method = lower_copy(value); + ok = method == "area" || method == "confidence"; + params->mask_filter_by_confidence = method == "confidence"; } else if (key == "mask_min_ratio") { ok = parse_float(value, ¶ms->mask_min_ratio); } else if (key == "mask_max_ratio") { @@ -805,6 +822,8 @@ static bool parse_adetailer_params(const sd_adetailer_params_t& public_params, ok = parse_int(value, ¶ms->mask_blur); } else if (key == "inpaint_padding") { ok = parse_int(value, ¶ms->inpaint_padding); + } else if (key == "inpaint_only_masked") { + ok = parse_bool(value, ¶ms->inpaint_only_masked); } else if (key == "inpaint_width") { ok = parse_int(value, ¶ms->inpaint_width); } else if (key == "inpaint_height") { @@ -898,14 +917,11 @@ bool adetail_image(adetailer_ctx_t* context, sd_image_t input_image, const sd_adetailer_params_t* public_params, const sd_img_gen_params_t* inpaint_params, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) { - *images_out = nullptr; - } - if (num_images_out != nullptr) { - *num_images_out = 0; - } + sd_image_generation_result_t** result_out) try { + if (result_out) + *result_out = nullptr; + if (!result_out) + return false; if (context == nullptr || context->detailer == nullptr || sd_ctx == nullptr || public_params == nullptr || inpaint_params == nullptr || input_image.data == nullptr) { return false; @@ -916,15 +932,19 @@ bool adetail_image(adetailer_ctx_t* context, return false; } - auto detections = context->detailer->predict(input_image, params); + auto resolved_inpaint = *inpaint_params; + resolved_inpaint.seed = sd::pipeline::resolve_seed(resolved_inpaint.seed); + const auto info = sd::adetailer_generation_info(*public_params, resolved_inpaint); + auto detections = context->detailer->predict(input_image, params); filter_and_sort_detections(&detections, input_image.width, input_image.height, params); auto masks = make_masks(detections, input_image.width, input_image.height, params); - sd_image_t current = copy_as_rgb(input_image); + sd_image_t current = copy_for_detailing(input_image); if (current.data == nullptr) { return false; } + std::unique_ptr current_owner(current.data, &free); int applied = 0; for (size_t i = 0; i < masks.size(); ++i) { std::string prompt = resolve_prompt(params.prompt, inpaint_params->prompt, static_cast(i)); @@ -939,22 +959,22 @@ bool adetail_image(adetailer_ctx_t* context, if (!mask_bbox(masks[i], &crop)) { continue; } - crop = expand_crop(crop, - current.width, - current.height, - params.inpaint_padding, - params.inpaint_width, - params.inpaint_height); + crop = params.inpaint_only_masked ? expand_crop(crop, + current.width, + current.height, + params.inpaint_padding, + params.inpaint_width, + params.inpaint_height) + : CropRegion{0, 0, static_cast(current.width), static_cast(current.height)}; sd_image_t local_image = resize_crop_image(current, crop, params.inpaint_width, params.inpaint_height); sd_image_t local_mask = resize_crop_mask(masks[i], crop, params.inpaint_width, params.inpaint_height); if (local_image.data == nullptr || local_mask.data == nullptr) { free(local_image.data); free(local_mask.data); - free(current.data); return false; } - sd_img_gen_params_t generation = *inpaint_params; + sd_img_gen_params_t generation = resolved_inpaint; generation.prompt = prompt.c_str(); generation.negative_prompt = negative_prompt.c_str(); generation.init_image = local_image; @@ -962,7 +982,7 @@ bool adetail_image(adetailer_ctx_t* context, generation.width = params.inpaint_width; generation.height = params.inpaint_height; generation.strength = params.denoising_strength; - generation.seed = inpaint_params->seed + static_cast(i); + generation.seed = static_cast((static_cast(resolved_inpaint.seed) + i) & INT64_MAX); generation.batch_count = 1; generation.ref_images = nullptr; generation.ref_images_count = 0; @@ -986,37 +1006,32 @@ bool adetail_image(adetailer_ctx_t* context, generation.sample_params.scheduler = params.scheduler; } - sd_image_t* generated = nullptr; - int generated_count = 0; - bool success = generate_image(sd_ctx, &generation, &generated, &generated_count); + sd_image_generation_result_t* generated = nullptr; + bool success = generate_image(sd_ctx, &generation, &generated); free(local_image.data); free(local_mask.data); - if (!success || generated_count <= 0 || generated == nullptr || generated[0].data == nullptr) { - free_sd_images(generated, generated_count); - free(current.data); + if (!success || generated == nullptr || generated->count <= 0 || generated->items[0].image.data == nullptr) { + free_sd_image_generation_result(generated); return false; } + std::unique_ptr generated_owner(generated, &free_sd_image_generation_result); Mask feather = gaussian_blur_mask(masks[i], params.mask_blur); - composite_crop(¤t, generated[0], feather, crop); - free_sd_images(generated, generated_count); + composite_crop(¤t, generated->items[0].image, feather, crop); ++applied; } + const std::vector infos{info}; auto* outputs = static_cast(calloc(1, sizeof(sd_image_t))); - if (outputs == nullptr) { - free(current.data); + if (!outputs) return false; - } outputs[0] = current; - if (images_out != nullptr) { - *images_out = outputs; - } else { - free_sd_images(outputs, 1); - } - if (num_images_out != nullptr) { - *num_images_out = 1; - } + current_owner.release(); + if (!sd::make_image_generation_result(outputs, 1, infos, result_out)) + return false; LOG_INFO("ADetailer applied %d mask(s)", applied); return true; +} catch (const std::exception& e) { + LOG_ERROR("ADetailer failed: %s", e.what()); + return false; } diff --git a/src/detailer.h b/src/detailer.h index 17a81a1dde..0e3ae50f5a 100644 --- a/src/detailer.h +++ b/src/detailer.h @@ -27,30 +27,32 @@ enum ADetailerSort { }; struct ADetailerParams { - const char* prompt = nullptr; - const char* negative_prompt = nullptr; - int input_size = 640; - float confidence = 0.3f; - float nms_threshold = 0.45f; - int max_detections = 100; - int mask_k_largest = 0; - float mask_min_ratio = 0.f; - float mask_max_ratio = 1.f; - int dilate_erode = 4; - int x_offset = 0; - int y_offset = 0; - bool merge_masks = false; - bool invert_mask = false; - int mask_blur = 4; - int inpaint_padding = 32; - int inpaint_width = 512; - int inpaint_height = 512; - float denoising_strength = 0.4f; - int steps = 0; - float cfg_scale = -1.f; - sample_method_t sample_method = SAMPLE_METHOD_COUNT; - scheduler_t scheduler = SCHEDULER_COUNT; - ADetailerSort sort_by = ADETAILER_SORT_NONE; + const char* prompt = nullptr; + const char* negative_prompt = nullptr; + int input_size = 640; + float confidence = 0.3f; + float nms_threshold = 0.45f; + int max_detections = 100; + int mask_k_largest = 0; + bool mask_filter_by_confidence = false; + float mask_min_ratio = 0.f; + float mask_max_ratio = 1.f; + int dilate_erode = 4; + int x_offset = 0; + int y_offset = 0; + bool merge_masks = false; + bool invert_mask = false; + int mask_blur = 4; + int inpaint_padding = 32; + bool inpaint_only_masked = true; + int inpaint_width = 512; + int inpaint_height = 512; + float denoising_strength = 0.4f; + int steps = 0; + float cfg_scale = -1.f; + sample_method_t sample_method = SAMPLE_METHOD_COUNT; + scheduler_t scheduler = SCHEDULER_COUNT; + ADetailerSort sort_by = ADETAILER_SORT_NONE; }; struct ADetailerGGML { diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index 0be0d2ee42..f2a0030cb1 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -66,7 +66,13 @@ struct DeviceResidencyManager { const std::vector& protected_params) = 0; virtual WeightPrefetchResult prefetch_params( uintptr_t owner_id, - const std::vector& tensors) = 0; + const std::vector& tensors) = 0; + virtual WeightPrefetchResult prepare_prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) { + return prefetch_params(owner_id, tensors); + } + virtual void submit_prefetched_params(uintptr_t owner_id) {} virtual bool activate_prefetched_params(uintptr_t owner_id, const std::vector& tensors) = 0; virtual void clear_prefetched_params(uintptr_t owner_id) = 0; diff --git a/src/extensions/photomaker_extension.cpp b/src/extensions/photomaker_extension.cpp index f904c8f731..59b9a2c2eb 100644 --- a/src/extensions/photomaker_extension.cpp +++ b/src/extensions/photomaker_extension.cpp @@ -154,12 +154,12 @@ struct PhotoMakerExtension : public GenerationExtension { PMVersion pm_version = std::strstr(model_path.c_str(), "v2") != nullptr ? PM_VERSION_2 : PM_VERSION_1; pmid_model = std::make_shared(ctx.backend_for(SDBackendModule::PHOTOMAKER), - ctx.tensor_storage_map, - "pmid", - ctx.version, - pm_version, - 20.f, - ctx.model_manager); + ctx.tensor_storage_map, + "pmid", + ctx.version, + pm_version, + 20.f, + ctx.model_manager); pmid_model->set_scale_overrides(ctx.params->linear_scale, ctx.params->attn_scale); if (pm_version == PM_VERSION_2) { LOG_INFO("using PhotoMaker Version 2"); diff --git a/src/generation_info.cpp b/src/generation_info.cpp new file mode 100644 index 0000000000..2f49054ea1 --- /dev/null +++ b/src/generation_info.cpp @@ -0,0 +1,276 @@ +#include "generation_info.h" + +#include +#include +#include + +#include "pipeline/request.h" +#include "runtime/guidance.h" + +namespace sd { + using nlohmann::json; + + static const char* text(const char* value) { + return value ? value : ""; + } + + static json sampling_info(const sd_sample_params_t& p, bool high_noise = false) { + const auto& g = p.guidance; + json info = {{"sample_steps", p.sample_steps}, {"guidance", {{"txt_cfg", g.txt_cfg}, {"distilled_guidance", g.distilled_guidance}}}}; + if (!high_noise && p.shifted_timestep > 0) + info["shifted_timestep"] = p.shifted_timestep; + if (std::isfinite(g.img_cfg) && g.img_cfg != 1.f) + info["guidance"]["img_cfg"] = g.img_cfg; + if (p.sample_method != SAMPLE_METHOD_COUNT) + info["sample_method"] = sd_sample_method_name(p.sample_method); + if (!high_noise && p.scheduler != SCHEDULER_COUNT) + info["scheduler"] = sd_scheduler_name(p.scheduler); + if (std::isfinite(p.eta)) + info["eta"] = p.eta; + if (!high_noise && std::isfinite(p.flow_shift)) + info["flow_shift"] = p.flow_shift; + if (!high_noise && p.custom_sigmas && p.custom_sigmas_count > 0) + info["custom_sigmas"] = std::vector(p.custom_sigmas, p.custom_sigmas + p.custom_sigmas_count); + if (p.extra_sample_args && *p.extra_sample_args) + info["extra_sample_args"] = p.extra_sample_args; + if (g.slg.layers && g.slg.layer_count > 0 && + (g.slg.scale != 0.f || guidance::parse_skip_layer_guidance_uncond_arg(p.extra_sample_args))) + info["guidance"]["slg"] = {{"scale", g.slg.scale}, {"layer_start", g.slg.layer_start}, {"layer_end", g.slg.layer_end}, {"layers", std::vector(g.slg.layers, g.slg.layers + g.slg.layer_count)}}; + return info; + } + + static json hires_info(const sd_hires_params_t& p) { + json info = {{"enabled", p.enabled}}; + if (!p.enabled) + return info; + info.update({{"upscaler", sd_hires_upscaler_name(p.upscaler)}, {"scale", p.scale}, {"target_width", p.target_width}, {"target_height", p.target_height}, {"steps", p.steps}, {"denoising_strength", p.denoising_strength}}); + if (p.upscaler == SD_HIRES_UPSCALER_MODEL) { + info["model_path"] = text(p.model_path); + info["upscale_tile_size"] = p.upscale_tile_size; + } + if (p.custom_sigmas && p.custom_sigmas_count > 0) + info["custom_sigmas"] = std::vector(p.custom_sigmas, p.custom_sigmas + p.custom_sigmas_count); + return info; + } + + static json cache_info(const sd_cache_params_t& p) { + static const char* modes[] = {"disabled", "easycache", "ucache", "dbcache", "taylorseer", "cache_dit", "spectrum"}; + if (p.mode < SD_CACHE_DISABLED || p.mode > SD_CACHE_SPECTRUM) + return {{"mode", static_cast(p.mode)}}; + json info = {{"mode", modes[p.mode]}}; + if (p.mode == SD_CACHE_DISABLED) + return info; + if (p.mode == SD_CACHE_EASYCACHE || p.mode == SD_CACHE_UCACHE) { + if (std::isfinite(p.reuse_threshold)) + info["reuse_threshold"] = p.reuse_threshold; + info.update({{"start_percent", p.start_percent}, {"end_percent", p.end_percent}}); + if (p.mode == SD_CACHE_UCACHE) + info.update({{"error_decay_rate", p.error_decay_rate}, {"use_relative_threshold", p.use_relative_threshold}, {"reset_error_on_compute", p.reset_error_on_compute}}); + } else if (p.mode == SD_CACHE_DBCACHE || p.mode == SD_CACHE_TAYLORSEER || p.mode == SD_CACHE_CACHE_DIT) { + // TaylorSeer also uses the shared DBCache thresholds and limits. + info.update({{"Fn_compute_blocks", p.Fn_compute_blocks}, + {"Bn_compute_blocks", p.Bn_compute_blocks}, + {"residual_diff_threshold", p.residual_diff_threshold}, + {"max_warmup_steps", p.max_warmup_steps}, + {"max_cached_steps", p.max_cached_steps}, + {"max_continuous_cached_steps", p.max_continuous_cached_steps}}); + if (p.mode == SD_CACHE_TAYLORSEER || p.mode == SD_CACHE_CACHE_DIT) + info.update({{"taylorseer_n_derivatives", p.taylorseer_n_derivatives}, {"taylorseer_skip_interval", p.taylorseer_skip_interval}}); + if (p.scm_mask && *p.scm_mask) + info.update({{"scm_mask", p.scm_mask}, {"scm_policy_dynamic", p.scm_policy_dynamic}}); + } else if (p.mode == SD_CACHE_SPECTRUM) { + info.update({{"spectrum_w", p.spectrum_w}, + {"spectrum_m", p.spectrum_m}, + {"spectrum_lam", p.spectrum_lam}, + {"spectrum_window_size", p.spectrum_window_size}, + {"spectrum_flex_window", p.spectrum_flex_window}, + {"spectrum_warmup_steps", p.spectrum_warmup_steps}, + {"spectrum_stop_percent", p.spectrum_stop_percent}}); + } + return info; + } + + template + static json common_info(const Params& p, const char* kind) { + json info = {{"schema", "sdcpp.generation/v1"}, {"kind", kind}, {"prompt", text(p.prompt)}, {"negative_prompt", text(p.negative_prompt)}, {"width", p.width}, {"height", p.height}, {"clip_skip", p.clip_skip}, {"seed", std::to_string(p.seed)}, {"sample_params", sampling_info(p.sample_params)}}; + if (p.hires.enabled) + info["hires"] = hires_info(p.hires); + if (p.cache.mode != SD_CACHE_DISABLED) + info["cache"] = cache_info(p.cache); + if (p.circular_x) + info["circular_x"] = true; + if (p.circular_y) + info["circular_y"] = true; + if (p.ref_image_args && *p.ref_image_args) + info["ref_image_args"] = p.ref_image_args; + if (p.image_preprocess.rules && *p.image_preprocess.rules) + info["image_preprocess"] = {{"rules", p.image_preprocess.rules}}; + const auto& t = p.vae_tiling_params; + if (t.enabled || t.temporal_tiling) { + auto& tiling = info["vae_tiling_params"]; + if (t.enabled) + tiling.update({{"enabled", true}, {"tile_size_w", t.tile_size_w}, {"tile_size_h", t.tile_size_h}, {"target_overlap", t.target_overlap}, {"rel_size_w", t.rel_size_w}, {"rel_size_h", t.rel_size_h}}); + if (t.temporal_tiling) + tiling["temporal_tiling"] = true; + if (t.extra_tiling_args && *t.extra_tiling_args) + tiling["extra_tiling_args"] = t.extra_tiling_args; + } + if (p.lora_count > 0) { + info["loras"] = json::array(); + for (uint32_t i = 0; i < p.lora_count; ++i) + info["loras"].push_back({{"path", text(p.loras[i].path)}, {"multiplier", p.loras[i].multiplier}, {"is_high_noise", p.loras[i].is_high_noise}}); + } + return info; + } + + static json image_info(const sd_img_gen_params_t& p) { + auto info = common_info(p, "image"); + if (p.init_image.data) + info["strength"] = p.strength; + if (p.control_image.data) + info["control_strength"] = p.control_strength; + if (p.ip_adapter_image.data) + info["ip_adapter_strength"] = p.ip_adapter_strength; + if (p.pm_params.id_images && p.pm_params.id_images_count > 0) { + info["pm_params"] = {{"style_strength", p.pm_params.style_strength}}; + if (p.pm_params.id_embed_path && *p.pm_params.id_embed_path) + info["pm_params"]["id_embed_path"] = p.pm_params.id_embed_path; + } + if (p.pulid_params.id_embedding_path && *p.pulid_params.id_embedding_path) + info["pulid_params"] = {{"id_embedding_path", p.pulid_params.id_embedding_path}, {"id_weight", p.pulid_params.id_weight}}; + return info; + } + + template + static Params resolved_params(const pipeline::GenerationRequest& request, Params p) { + p.width = request.width; + p.height = request.height; + p.seed = request.seed; + p.clip_skip = request.clip_skip; + p.sample_params = request.sample_params; + p.hires = request.hires; + p.vae_tiling_params = request.vae_tiling_params; + return p; + } + + json generation_info(const pipeline::GenerationRequest& request, const sd_img_gen_params_t& params) { + auto p = resolved_params(request, params); + p.qwen_image_layers = request.qwen_image_layers; + auto info = image_info(p); + if (!request.models.empty()) + info["models"] = request.models; + if (request.is_layered_image) + info["qwen_image_layers"] = p.qwen_image_layers; + return info; + } + + json generation_info(const pipeline::GenerationRequest& request, const sd_vid_gen_params_t& params) { + const auto p = resolved_params(request, params); + auto info = common_info(p, "video"); + if (!request.models.empty()) + info["models"] = request.models; + info.update({{"video_frames", request.frames}, {"fps", request.fps}}); + if (p.init_image.data || p.end_image.data) + info["strength"] = p.strength; + if (request.use_vace) + info["vace_strength"] = p.vace_strength; + if (request.use_high_noise) + info.update({{"moe_boundary", p.moe_boundary}, {"high_noise_sample_params", sampling_info(request.high_noise_sample_params, true)}}); + if (p.ref_videos_count > 0) { + info["ref_videos"] = json::array(); + for (int i = 0; i < p.ref_videos_count; ++i) + info["ref_videos"].push_back({{"fps", p.ref_videos[i].fps}}); + } + return info; + } + + json adetailer_generation_info(const sd_adetailer_params_t& params, const sd_img_gen_params_t& inpaint_params) { + auto inpaint = image_info(inpaint_params); + inpaint["strength"] = inpaint_params.strength; + // ADetailer clears these inputs before generating each crop. + for (const auto* key : {"control_strength", "pm_params", "pulid_params", "hires", "image_preprocess"}) + inpaint.erase(key); + sd_img_gen_params_t defaults; + sd_img_gen_params_init(&defaults); + if (inpaint_params.qwen_image_layers != defaults.qwen_image_layers) + inpaint["qwen_image_layers"] = inpaint_params.qwen_image_layers; + json adetailer = json::object(); + if (params.prompt && *params.prompt) + adetailer["prompt"] = params.prompt; + if (params.negative_prompt && *params.negative_prompt) + adetailer["negative_prompt"] = params.negative_prompt; + if (params.extra_ad_args && *params.extra_ad_args) + adetailer["extra_ad_args"] = params.extra_ad_args; + return {{"schema", "sdcpp.generation/v1"}, {"kind", "image"}, {"adetailer", adetailer}, {"inpaint_params", inpaint}}; + } + + json upscale_generation_info(uint32_t upscale_factor) { + return {{"schema", "sdcpp.generation/v1"}, {"kind", "image"}, {"upscale", {{"upscale_factor", upscale_factor}}}}; + } + + char* copy_generation_info(const json& info) { + const std::string text = info.dump(-1, ' ', false, json::error_handler_t::replace); + char* result = static_cast(malloc(text.size() + 1)); + if (result) + memcpy(result, text.c_str(), text.size() + 1); + return result; + } + + bool make_image_generation_result(sd_image_t* images, int count, const std::vector& infos, sd_image_generation_result_t** result_out) { + if (!result_out || !images || count <= 0 || infos.size() != static_cast(count)) { + free_sd_images(images, count); + return false; + } + auto* result = static_cast(calloc(1, sizeof(sd_image_generation_result_t))); + if (!result) { + free_sd_images(images, count); + return false; + } + result->items = static_cast(calloc(count, sizeof(sd_image_result_t))); + if (!result->items) { + free(result); + free_sd_images(images, count); + return false; + } + result->count = count; + try { + for (int i = 0; i < count; ++i) { + result->items[i].image = images[i]; + images[i] = {}; + result->items[i].generation_info = copy_generation_info(infos[i]); + if (!result->items[i].generation_info) { + free_sd_images(images, count); + free_sd_image_generation_result(result); + return false; + } + } + } catch (...) { + free_sd_images(images, count); + free_sd_image_generation_result(result); + return false; + } + free(images); + *result_out = result; + return true; + } +} + +void free_sd_image_generation_result(sd_image_generation_result_t* result) { + if (!result) + return; + for (int i = 0; i < result->count; ++i) { + free(result->items[i].image.data); + free(result->items[i].generation_info); + } + free(result->items); + free(result); +} + +void free_sd_video_generation_result(sd_video_generation_result_t* result) { + if (!result) + return; + free_sd_images(result->frames, result->frame_count); + free_sd_audio(result->audio); + free(result->generation_info); + free(result); +} diff --git a/src/generation_info.h b/src/generation_info.h new file mode 100644 index 0000000000..56688ee812 --- /dev/null +++ b/src/generation_info.h @@ -0,0 +1,23 @@ +#ifndef __SD_GENERATION_INFO_H__ +#define __SD_GENERATION_INFO_H__ + +#include + +#include "json.hpp" +#include "stable-diffusion.h" + +namespace sd::pipeline { + struct GenerationRequest; +} + +namespace sd { + nlohmann::json generation_info(const pipeline::GenerationRequest& request, const sd_img_gen_params_t& params); + nlohmann::json generation_info(const pipeline::GenerationRequest& request, const sd_vid_gen_params_t& params); + nlohmann::json adetailer_generation_info(const sd_adetailer_params_t& params, const sd_img_gen_params_t& inpaint_params); + nlohmann::json upscale_generation_info(uint32_t upscale_factor); + char* copy_generation_info(const nlohmann::json& info); + // Consumes images even on failure. + bool make_image_generation_result(sd_image_t* images, int count, const std::vector& infos, sd_image_generation_result_t** result_out); +} + +#endif // __SD_GENERATION_INFO_H__ diff --git a/src/model.h b/src/model.h index 0c1a37a66c..4695f5c74d 100644 --- a/src/model.h +++ b/src/model.h @@ -68,9 +68,36 @@ enum SDVersion { VERSION_IRIS, VERSION_FLUX3, VERSION_MOSS_IMAGE, + VERSION_FIBO2, + VERSION_CHROMA, + VERSION_P_H3_IMAGE, VERSION_COUNT, }; +enum class ModelVAEKind { KL, + FLUX, + FLUX2, + WAN, + HUNYUAN, + LTXAV, + MINIMAX_H3, + MAGE }; + +struct ModelTraits { + bool is_dit = false; + bool supports_video_generation = false; + bool is_inpaint = false; + bool is_unet_edit = false; + bool is_control = false; + ModelVAEKind vae_kind = ModelVAEKind::KL; + bool video_tae = false; + bool fake_vae = false; + bool vae_format_override = false; +}; + +const ModelTraits* find_model_traits(SDVersion version); +const char* sd_version_name(SDVersion version); + static inline bool sd_version_is_sd1(SDVersion version) { if (version == VERSION_SD1 || version == VERSION_SD1_INPAINT || version == VERSION_SD1_PIX2PIX || version == VERSION_SD1_TINY_UNET || version == VERSION_SDXS_512_DS) { return true; @@ -114,6 +141,7 @@ static inline bool sd_version_is_flux(SDVersion version) { version == VERSION_FLUX_CONTROLS || version == VERSION_FLEX_2 || version == VERSION_OVIS_IMAGE || + version == VERSION_CHROMA || version == VERSION_CHROMA_RADIANCE) { return true; } @@ -139,7 +167,7 @@ static inline bool sd_version_is_ltxav(SDVersion version) { } static inline bool sd_version_is_minimax_h3(SDVersion version) { - return version == VERSION_MINIMAX_H3; + return version == VERSION_MINIMAX_H3 || version == VERSION_P_H3_IMAGE; } static inline bool sd_version_is_wan(SDVersion version) { @@ -271,7 +299,8 @@ static inline bool sd_version_is_pixart(SDVersion version) { } static inline bool sd_version_supports_video_generation(SDVersion version) { - return version == VERSION_SVD || sd_version_is_wan(version) || sd_version_is_hunyuan_video(version) || sd_version_is_lingbot_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version); + const auto* traits = find_model_traits(version); + return traits && traits->supports_video_generation; } static inline bool sd_version_supports_image_generation(SDVersion version) { @@ -279,113 +308,43 @@ static inline bool sd_version_supports_image_generation(SDVersion version) { } static inline bool sd_version_uses_flux_vae(SDVersion version) { - if (sd_version_is_flux(version) || sd_version_is_z_image(version) || sd_version_is_boogu_image(version) || sd_version_is_longcat(version)) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::FLUX; } static inline bool sd_version_uses_flux2_vae(SDVersion version) { - if (sd_version_is_flux2(version) || - sd_version_is_flux3(version) || - sd_version_is_ernie_image(version) || - sd_version_is_lens(version) || - sd_version_is_ideogram4(version) || - sd_version_is_sefi_image(version) || - sd_version_is_llada_image(version) || - version == VERSION_MOSS_IMAGE) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::FLUX2; } static inline bool sd_version_uses_wan_vae(SDVersion version) { - if (sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_qwen_image(version) || sd_version_is_krea2(version) || sd_version_is_anima(version) || version == VERSION_MING_IMAGE) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::WAN; } static inline bool sd_version_uses_hunyuan_video_vae(SDVersion version) { - return sd_version_is_hunyuan_video(version); + const auto* traits = find_model_traits(version); + return traits && traits->vae_kind == ModelVAEKind::HUNYUAN; } static inline bool sd_version_is_inpaint(SDVersion version) { - if (version == VERSION_SD1_INPAINT || - version == VERSION_SD2_INPAINT || - version == VERSION_SDXL_INPAINT || - version == VERSION_FLUX_FILL || - version == VERSION_FLEX_2) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->is_inpaint; } static inline bool sd_version_is_dit(SDVersion version) { - if (sd_version_is_flux(version) || - sd_version_is_flux2(version) || - sd_version_is_flux3(version) || - sd_version_is_ltxav(version) || - sd_version_is_minimax_h3(version) || - sd_version_is_sd3(version) || - sd_version_is_wan(version) || - sd_version_is_lingbot_video(version) || - sd_version_is_qwen_image(version) || - sd_version_is_hunyuan_video(version) || - version == VERSION_HIDREAM_O1 || - sd_version_is_anima(version) || - sd_version_is_z_image(version) || - sd_version_is_z_image_l2p(version) || - version == VERSION_MING_IMAGE || - sd_version_is_llada_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_ernie_image(version) || - sd_version_is_lens(version) || - sd_version_is_minit2i(version) || - sd_version_is_longcat(version) || - sd_version_is_pid(version) || - sd_version_is_ideogram4(version) || - sd_version_is_sefi_image(version) || - sd_version_is_krea2(version) || - sd_version_is_mage_flow(version) || - sd_version_is_sensenova_u1(version) || - sd_version_is_pixart(version) || - version == VERSION_IRIS || - version == VERSION_MOSS_IMAGE) { - return true; - } - return false; + const auto* traits = find_model_traits(version); + return traits && traits->is_dit; } static inline bool sd_version_is_unet_edit(SDVersion version) { - return version == VERSION_SD1_PIX2PIX || version == VERSION_SDXL_PIX2PIX; + const auto* traits = find_model_traits(version); + return traits && traits->is_unet_edit; } static inline bool sd_version_is_control(SDVersion version) { - return version == VERSION_FLUX_CONTROLS || version == VERSION_FLEX_2; -} - -// These model families consume uploaded reference images as generation -// conditioning. Inpaint and control variants keep their init-image semantics. -static inline bool sd_version_supports_ref_image_conditioning(SDVersion version) { - if (sd_version_is_inpaint(version) || sd_version_is_control(version) || version == VERSION_CHROMA_RADIANCE) { - return false; - } - - return sd_version_is_unet_edit(version) || - sd_version_is_flux(version) || - sd_version_is_flux3(version) || - sd_version_is_flux2(version) || - sd_version_is_qwen_image(version) || - sd_version_is_longcat(version) || - sd_version_is_z_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_llada_image(version) || - sd_version_is_krea2(version) || - sd_version_is_mage_flow(version) || - sd_version_is_anima(version) || - sd_version_is_sefi_image(version) || - sd_version_is_pid(version); + const auto* traits = find_model_traits(version); + return traits && traits->is_control; } static bool sd_version_is_inpaint_or_unet_edit(SDVersion version) { diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 1e5377ccc8..b2657a3691 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -69,7 +69,8 @@ struct LoraModel : public GGMLRunner { metadata_alpha = -1.f; const auto metadata = model_manager->loader().file_metadata(source_file); const auto alpha = metadata.find("alpha"); - if (alpha != metadata.end()) { + // Converters that write "baked_scale" have already folded alpha / rank into lora_up. + if (alpha != metadata.end() && metadata.find("baked_scale") == metadata.end()) { try { size_t parsed = 0; metadata_alpha = std::stof(alpha->second, &parsed); diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp index 5a5e6615cc..e79ee2f12a 100644 --- a/src/model/common/block.hpp +++ b/src/model/common/block.hpp @@ -197,16 +197,12 @@ class GEGLU : public UnaryBlock { // return: [ne3, ne2, ne1, dim_out] auto proj = std::dynamic_pointer_cast(blocks["proj"]); - x = proj->forward(ctx, x); // [ne3, ne2, ne1, dim_out*2] - auto x_vec = ggml_ext_chunk(ctx->ggml_ctx, x, 2, 0, false); - x = x_vec[0]; // [ne3, ne2, ne1, dim_out] - auto gate = x_vec[1]; // [ne3, ne2, ne1, dim_out] + x = proj->forward(ctx, x); // [ne3, ne2, ne1, dim_out*2] - gate = ggml_cont(ctx->ggml_ctx, gate); - - gate = ggml_ext_gelu(ctx->ggml_ctx, gate, true); - - x = ggml_mul(ctx->ggml_ctx, x, gate); // [ne3, ne2, ne1, dim_out] + if (!ggml_is_contiguous_1(x)) { + x = ggml_cont(ctx->ggml_ctx, x); + } + x = ggml_geglu_swapped(ctx->ggml_ctx, x); // [ne3, ne2, ne1, dim_out] return x; } diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index a0b342240d..0a888e020a 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -661,7 +661,7 @@ namespace Rope { const int axes_dim_num = static_cast(axes_dim.size()); auto make_image_ids = [=, &result](int image_h, int image_w, int image_index) { std::vector> image_ids(static_cast(bs) * image_h * image_w, - std::vector(axes_dim_num, 0.f)); + std::vector(axes_dim_num, 0.f)); result.positions.append_image(image_h, image_w); int h_start = -(image_h - image_h / 2); int w_start = -(image_w - image_w / 2); @@ -1078,8 +1078,8 @@ namespace Rope { if (backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32 && x->ne[0] % 2 == 0 && pe->ne[0] == 2 && pe->ne[1] == 2 && pe->ne[2] == x->ne[0] / 2 && pe->ne[3] == x->ne[2]) { - auto out = ggml_rope_apply(ctx, x, pe, rope_interleaved); - if (ggml_backend_supports_op(backend, out)) { + auto out = sd_ggml_rope_apply(ctx, backend, x, pe, rope_interleaved); + if (out && ggml_backend_supports_op(backend, out)) { return out; } } diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index a414d96e63..e9e33b88a1 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -7,6 +7,7 @@ #include #include #include +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -128,7 +129,7 @@ namespace Anima { auto gate = emb_chunks[2]; auto x = norm->forward(ctx, hidden_states); - x = Flux::modulate(ctx->ggml_ctx, x, shift, scale); + x = Flux::modulate(ctx, x, shift, scale); return {x, gate}; } @@ -168,7 +169,7 @@ namespace Anima { auto scale = emb_chunks[1]; auto x = norm->forward(ctx, hidden_states); - x = Flux::modulate(ctx->ggml_ctx, x, shift, scale); + x = Flux::modulate(ctx, x, shift, scale); return x; } }; @@ -562,10 +563,9 @@ namespace Anima { }; struct AnimaRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: - std::vector image_pe_vec; - std::vector adapter_q_pe_vec; - std::vector adapter_k_pe_vec; AnimaConfig config; AnimaNet net; @@ -660,19 +660,23 @@ namespace Anima { int64_t h_pad = x->ne[1] + pad_h; int64_t w_pad = x->ne[0] + pad_w; - image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1, - static_cast(h_pad), - static_cast(w_pad), - static_cast(config.patch_size), - config.theta, - config.axes_dim, - 4.0f, - 4.0f, - 1.0f, - ref_latents)); - int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); - auto image_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, image_pos_len); - set_backend_tensor_data(image_pe, image_pe_vec.data()); + sd::PECacheKey image_pe_key("anima.image_pe"); + image_pe_key.add(h_pad).add(w_pad).add(config.patch_size).add(config.theta).add(config.axes_dim).add(config.head_dim); + image_pe_key.add_shapes(ref_latents).add(circular_x_enabled).add(circular_y_enabled); + auto image_pe = make_input(pe_cache_.get_or_create(image_pe_key, [&]() { + auto image_pe_vec = finish_rope_pe(gen_anima_image_pe_vec(1, + static_cast(h_pad), + static_cast(w_pad), + static_cast(config.patch_size), + config.theta, + config.axes_dim, + 4.0f, + 4.0f, + 1.0f, + ref_latents)); + int64_t image_pos_len = static_cast(image_pe_vec.size()) / (2 * 2 * (config.head_dim / 2)); + return sd::Tensor({2, 2, config.head_dim / 2, image_pos_len}, std::move(image_pe_vec)); + })); ggml_tensor* adapter_q_pe = nullptr; ggml_tensor* adapter_k_pe = nullptr; @@ -680,16 +684,16 @@ namespace Anima { int64_t target_len = t5_ids->ne[0]; int64_t source_len = context->ne[1]; - adapter_q_pe_vec = gen_1d_rope_pe_vec(target_len, 64, 10000.f); - adapter_k_pe_vec = gen_1d_rope_pe_vec(source_len, 64, 10000.f); - - int64_t target_pos_len = static_cast(adapter_q_pe_vec.size()) / (2 * 2 * 32); - int64_t source_pos_len = static_cast(adapter_k_pe_vec.size()) / (2 * 2 * 32); - - adapter_q_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, target_pos_len); - adapter_k_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 32, source_pos_len); - set_backend_tensor_data(adapter_q_pe, adapter_q_pe_vec.data()); - set_backend_tensor_data(adapter_k_pe, adapter_k_pe_vec.data()); + sd::PECacheKey q_key("anima.adapter_q_pe"); + q_key.add(target_len); + adapter_q_pe = make_input(pe_cache_.get_or_create(q_key, [&]() { + return sd::Tensor({2, 2, 32, target_len}, gen_1d_rope_pe_vec(target_len, 64, 10000.f)); + })); + sd::PECacheKey k_key("anima.adapter_k_pe"); + k_key.add(source_len); + adapter_k_pe = make_input(pe_cache_.get_or_create(k_key, [&]() { + return sd::Tensor({2, 2, 32, source_len}, gen_1d_rope_pe_vec(source_len, 64, 10000.f)); + })); } auto runner_ctx = get_context(); diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 1ca1370984..3f652e1589 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -6,6 +6,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -127,9 +128,9 @@ namespace Boogu { } }; - __STATIC_INLINE__ ggml_tensor* scale_modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) { - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); - return ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + __STATIC_INLINE__ ggml_tensor* scale_modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) { + scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); + return ggml_ext_modulate(ctx, x, nullptr, scale); } __STATIC_INLINE__ ggml_tensor* gate_residual(ggml_context* ctx, ggml_tensor* residual, ggml_tensor* x, ggml_tensor* gate) { @@ -185,7 +186,7 @@ namespace Boogu { auto scale_mlp = mods[2]; auto gate_mlp = mods[3]; - x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), scale_msa); + x = scale_modulate(ctx, norm->forward(ctx, x), scale_msa); return {x, gate_msa, scale_mlp, gate_mlp}; } }; @@ -230,7 +231,7 @@ namespace Boogu { auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); auto emb = linear_1->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning_embedding)); - x = scale_modulate(ctx->ggml_ctx, norm->forward(ctx, x), emb); + x = scale_modulate(ctx, norm->forward(ctx, x), emb); x = linear_2->forward(ctx, x); return x; } @@ -333,7 +334,7 @@ namespace Boogu { auto attn_output = attn->forward(ctx, norm_hidden_states, norm_hidden_states, rotary_emb, attention_mask); hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, norm2->forward(ctx, attn_output), gate_msa); - auto mlp_input = scale_modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp); + auto mlp_input = scale_modulate(ctx, ffn_norm1->forward(ctx, hidden_states), scale_mlp); auto mlp_output = feed_forward->forward(ctx, mlp_input); hidden_states = gate_residual(ctx->ggml_ctx, hidden_states, ffn_norm2->forward(ctx, mlp_output), gate_mlp); } else { @@ -511,14 +512,14 @@ namespace Boogu { img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_attn_norm->forward(ctx, img_attn_out), img_gate_msa); img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_self_attn_norm->forward(ctx, img_self_attn_out), img_gate_self); - auto img_mlp_input = scale_modulate(ctx->ggml_ctx, img_norm2_out, img_scale_mlp); + auto img_mlp_input = scale_modulate(ctx, img_norm2_out, img_scale_mlp); img_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, img_shift_mlp, img_shift_mlp->ne[0], 1, img_shift_mlp->ne[1]); img_mlp_input = ggml_add(ctx->ggml_ctx, img_mlp_input, img_shift_mlp); auto img_mlp_out = img_feed_forward->forward(ctx, img_ffn_norm1->forward(ctx, img_mlp_input)); img_hidden_states = gate_residual(ctx->ggml_ctx, img_hidden_states, img_ffn_norm2->forward(ctx, img_mlp_out), img_gate_mlp); instruct_hidden_states = gate_residual(ctx->ggml_ctx, instruct_hidden_states, instruct_attn_norm->forward(ctx, instruct_attn_out), instruct_gate_msa); - auto instruct_mlp_input = scale_modulate(ctx->ggml_ctx, instruct_norm2_out, instruct_scale_mlp); + auto instruct_mlp_input = scale_modulate(ctx, instruct_norm2_out, instruct_scale_mlp); instruct_shift_mlp = ggml_reshape_3d(ctx->ggml_ctx, instruct_shift_mlp, instruct_shift_mlp->ne[0], 1, instruct_shift_mlp->ne[1]); instruct_mlp_input = ggml_add(ctx->ggml_ctx, instruct_mlp_input, instruct_shift_mlp); auto instruct_mlp_out = instruct_feed_forward->forward(ctx, instruct_ffn_norm1->forward(ctx, instruct_mlp_input)); @@ -543,18 +544,18 @@ namespace Boogu { blocks["x_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true); blocks["ref_image_patch_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, this->config.hidden_size, true); blocks["time_caption_embed"] = std::make_shared(this->config.hidden_size, - this->config.instruction_feat_dim, - 256, - this->config.norm_eps, - this->config.timestep_scale); + this->config.instruction_feat_dim, + 256, + this->config.norm_eps, + this->config.timestep_scale); for (int i = 0; i < this->config.num_refiner_layers; i++) { blocks["noise_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, - this->config.num_attention_heads, - this->config.num_kv_heads, - this->config.multiple_of, - this->config.norm_eps, - true); + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + true); blocks["ref_image_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, this->config.num_attention_heads, this->config.num_kv_heads, @@ -562,11 +563,11 @@ namespace Boogu { this->config.norm_eps, true); blocks["context_refiner." + std::to_string(i)] = std::make_shared(this->config.hidden_size, - this->config.num_attention_heads, - this->config.num_kv_heads, - this->config.multiple_of, - this->config.norm_eps, - false); + this->config.num_attention_heads, + this->config.num_kv_heads, + this->config.multiple_of, + this->config.norm_eps, + false); } for (int i = 0; i < this->config.num_double_stream_layers; i++) { @@ -759,9 +760,9 @@ namespace Boogu { } struct BooguImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; BooguConfig config; BooguImageModel boogu; - std::vector pe_vec; BooguImageRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -799,17 +800,21 @@ namespace Boogu { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("boogu.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(gen_boogu_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = boogu.forward(&runner_ctx, x, timesteps, context, pe, ref_latents); diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index b284cb83f7..15badb9834 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -4,6 +4,7 @@ #include #include #include +#include "core/pe_cache.h" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" @@ -203,9 +204,8 @@ namespace ErnieImage { auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); auto gate = gate_proj->forward(ctx, x); - gate = ggml_ext_gelu(ctx->ggml_ctx, gate); x = up_proj->forward(ctx, x); - x = ggml_mul(ctx->ggml_ctx, x, gate); + x = ggml_geglu_split(ctx->ggml_ctx, gate, x); x = linear_fc2->forward(ctx, x); return x; } @@ -246,13 +246,13 @@ namespace ErnieImage { auto residual = x; x = adaLN_sa_ln->forward(ctx, x); - x = Flux::modulate(ctx->ggml_ctx, x, shift_msa, scale_msa, true); + x = Flux::modulate(ctx, x, shift_msa, scale_msa, true); auto attn_out = self_attention->forward(ctx, x, pe, attention_mask); x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa)); residual = x; x = adaLN_mlp_ln->forward(ctx, x); - x = Flux::modulate(ctx->ggml_ctx, x, shift_mlp, scale_mlp, true); + x = Flux::modulate(ctx, x, shift_mlp, scale_mlp, true); x = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, x), gate_mlp)); return x; } @@ -274,7 +274,7 @@ namespace ErnieImage { auto shift = mods[1]; x = norm->forward(ctx, x); - x = Flux::modulate(ctx->ggml_ctx, x, shift, scale); + x = Flux::modulate(ctx, x, shift, scale); return x; } }; @@ -383,9 +383,9 @@ namespace ErnieImage { }; struct ErnieImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; ErnieImageConfig config; ErnieImageModel ernie_image; - std::vector pe_vec; ErnieImageRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -415,16 +415,20 @@ namespace ErnieImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, config.axes_dim_sum, 1, pos_len, 2); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("ernie_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_ernie_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({config.axes_dim_sum, 1, pos_len, 2}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = ernie_image.forward(&runner_ctx, x, timesteps, context, pe); diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp new file mode 100644 index 0000000000..0c42235f8f --- /dev/null +++ b/src/model/diffusion/fibo2.hpp @@ -0,0 +1,403 @@ +#ifndef __SD_MODEL_DIFFUSION_FIBO2_HPP__ +#define __SD_MODEL_DIFFUSION_FIBO2_HPP__ + +#include "core/pe_cache.h" +#include "model/diffusion/z_image.hpp" + +// https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_bria_fibo2.py +namespace Fibo2 { + struct Fibo2Config { + int64_t dim = 3584; + int64_t channels = 128; + int64_t text_dim = 7680; + int64_t head_dim = 128; + int64_t depth = 44; + int64_t refiner_depth = 2; + int64_t perceiver_depth = 2; + int64_t max_gist = 512; + int64_t ffn_dim = 9536; + bool fused_qkv = false; + bool fused_kv = false; + bool fused_ffn = false; + std::vector injection_layers = {6, 14, 22, 30, 38}; + + static Fibo2Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) { + Fibo2Config config; + int64_t depth = 0, refiner_depth = 0, perceiver_depth = 0; + std::set injection_layers; + for (const auto& [name, tensor] : weights) { + if (!starts_with(name, prefix + ".")) { + continue; + } + auto key = name.substr(prefix.size() + 1); + if (key == "x_embedder.weight") { + config.channels = tensor.ne[0]; + config.dim = tensor.ne[1]; + } else if (key == "perceiver.latent_queries") { + config.max_gist = tensor.ne[1]; + } else if (key == "perceiver.layers.0.cross_attention.to_k.weight" || key == "perceiver.layers.0.cross_attention.to_kv.weight") { + config.text_dim = tensor.ne[0]; + config.fused_kv = key == "perceiver.layers.0.cross_attention.to_kv.weight"; + } else if (key == "layers.0.attention.norm_q.weight") { + config.head_dim = tensor.ne[0]; + } else if (key == "layers.0.attention.qkv.weight") { + config.fused_qkv = true; + } else if (key == "layers.0.feed_forward.w1.weight") { + config.ffn_dim = tensor.ne[1]; + } else if (key == "layers.0.feed_forward.w13.weight") { + config.ffn_dim = tensor.ne[1] / 2; + config.fused_ffn = true; + } + if (starts_with(key, "layers.")) { + int index = std::stoi(key.substr(7)); + depth = std::max(depth, index + 1); + if (ends_with(key, ".cross_attn_gate")) { + injection_layers.insert(index); + } + } else if (starts_with(key, "noise_refiner.")) { + refiner_depth = std::max(refiner_depth, std::stoi(key.substr(14)) + 1); + } else if (starts_with(key, "perceiver.layers.")) { + perceiver_depth = std::max(perceiver_depth, std::stoi(key.substr(17)) + 1); + } + } + if (depth) { + config.depth = depth; + config.refiner_depth = refiner_depth; + config.perceiver_depth = perceiver_depth; + config.injection_layers.assign(injection_layers.begin(), injection_layers.end()); + LOG_VERBOSE("fibo2: dim = %" PRId64 ", depth = %" PRId64 ", head_dim = %" PRId64 ", text_dim = %" PRId64, + config.dim, config.depth, config.head_dim, config.text_dim); + } + return config; + } + + int gist_count(int text_length) const { + double ratio = std::clamp((text_length - 256.0) / (2048.0 - 256.0), 0.0, 1.0); + return std::min(static_cast(max_gist), static_cast(std::ceil((256 + ratio * (max_gist - 256)) / 64)) * 64); + } + }; + + struct Attention : public GGMLBlock { + int64_t head_dim; + int64_t heads; + bool fused_qkv; + bool fused_kv; + + Attention(const Fibo2Config& config, int64_t kv_dim, bool joint = false) + : head_dim(config.head_dim), heads(config.dim / config.head_dim), fused_qkv(joint && config.fused_qkv), fused_kv(!joint && config.fused_kv) { + if (fused_qkv) { + blocks["qkv"] = std::make_shared(config.dim, 3 * config.dim, false); + } else { + blocks["to_q"] = std::make_shared(config.dim, config.dim, false); + if (fused_kv) { + blocks["to_kv"] = std::make_shared(kv_dim, 2 * config.dim, false); + } else { + blocks["to_k"] = std::make_shared(kv_dim, config.dim, false); + blocks["to_v"] = std::make_shared(kv_dim, config.dim, false); + } + } + blocks["to_out.0"] = std::make_shared(config.dim, config.dim, false); + blocks["norm_q"] = std::make_shared(head_dim, 1e-5f); + blocks["norm_k"] = std::make_shared(head_dim, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* kv, ggml_tensor* pe = nullptr) { + auto c = ctx->ggml_ctx; + auto project = [&](const char* name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + ggml_tensor *q, *k, *v; + if (fused_qkv) { + auto parts = ggml_ext_chunk(c, project("qkv", x), 3, 0); + q = parts[0]; + k = parts[1]; + v = parts[2]; + } else { + q = project("to_q", x); + if (fused_kv) { + auto parts = ggml_ext_chunk(c, project("to_kv", kv), 2, 0); + k = parts[0]; + v = parts[1]; + } else { + k = project("to_k", kv); + v = project("to_v", kv); + } + } + auto reshape = [&](ggml_tensor* h) { + return ggml_reshape_4d(c, h, head_dim, heads, h->ne[1], h->ne[2]); + }; + q = std::dynamic_pointer_cast(blocks["norm_q"])->forward(ctx, reshape(q)); + k = std::dynamic_pointer_cast(blocks["norm_k"])->forward(ctx, reshape(k)); + v = reshape(v); + auto prepare = [&](ggml_tensor* h) { + if (pe) { + return Rope::apply_rope(ctx, h, pe); + } + h = ggml_cont(c, ggml_permute(c, h, 0, 2, 1, 3)); + return ggml_reshape_3d(c, h, head_dim, h->ne[1], heads * h->ne[3]); + }; + auto y = ggml_ext_attention_ext(ctx, prepare(q), prepare(k), v, heads, nullptr, true, ctx->flash_attn_enabled); + return std::dynamic_pointer_cast(blocks["to_out.0"])->forward(ctx, y); + } + }; + + struct FeedForward : public ZImage::FeedForward { + bool fused; + + FeedForward(const Fibo2Config& config) + : ZImage::FeedForward(config.dim, config.ffn_dim, 1), fused(config.fused_ffn) { + if (fused) { + blocks.erase("w1"); + blocks.erase("w3"); + blocks["w13"] = std::make_shared(config.dim, 2 * config.ffn_dim, false); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (!fused) { + return ZImage::FeedForward::forward(ctx, x); + } + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { + w2->set_force_prec_f32(true); + } + x = std::dynamic_pointer_cast(blocks["w13"])->forward(ctx, x); + return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, x)); + } + }; + + struct PerceiverLayer : public GGMLBlock { + PerceiverLayer(const Fibo2Config& config) { + for (const char* name : {"cross_attn_norm", "self_attn_norm", "ffn_norm"}) { + blocks[name] = std::make_shared(config.dim, 1e-5f); + } + blocks["cross_attention"] = std::make_shared(config, config.text_dim); + blocks["self_attention"] = std::make_shared(config, config.dim); + blocks["ffn"] = std::make_shared(config); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* text) { + auto norm = [&](const char* name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + auto h = std::dynamic_pointer_cast(blocks["cross_attention"])->forward(ctx, norm("cross_attn_norm", x), text); + x = ggml_add(ctx->ggml_ctx, x, h); + h = norm("self_attn_norm", x); + h = std::dynamic_pointer_cast(blocks["self_attention"])->forward(ctx, h, h); + x = ggml_add(ctx->ggml_ctx, x, h); + h = std::dynamic_pointer_cast(blocks["ffn"])->forward(ctx, norm("ffn_norm", x)); + return ggml_add(ctx->ggml_ctx, x, h); + } + }; + + struct Perceiver : public GGMLBlock { + Fibo2Config config; + + Perceiver(const Fibo2Config& config) + : config(config) { + for (int i = 0; i < config.perceiver_depth; ++i) { + blocks["layers." + std::to_string(i)] = std::make_shared(config); + } + blocks["norm_out"] = std::make_shared(config.dim, 1e-5f); + } + + void init_params(ggml_context* ctx, const String2TensorStorage& weights = {}, const std::string prefix = "") override { + params["latent_queries"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config.dim, config.max_gist); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* text, int gist_count) { + auto x = ggml_ext_slice(ctx->ggml_ctx, params["latent_queries"], 1, 0, gist_count); + for (int i = 0; i < config.perceiver_depth; ++i) { + x = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, x, text); + } + return std::dynamic_pointer_cast(blocks["norm_out"])->forward(ctx, x); + } + }; + + struct TransformerBlock : public GGMLBlock { + bool injection; + + TransformerBlock(const Fibo2Config& config, bool injection) + : injection(injection) { + blocks["attention"] = std::make_shared(config, config.dim, true); + blocks["feed_forward"] = std::make_shared(config); + for (const char* name : {"attention_norm1", "attention_norm2", "ffn_norm1", "ffn_norm2"}) { + blocks[name] = std::make_shared(config.dim, 1e-5f); + } + blocks["adaLN_modulation.0"] = std::make_shared(std::min(config.dim, 256), 4 * config.dim); + if (injection) { + blocks["cross_attention"] = std::make_shared(config, config.text_dim); + blocks["cross_attn_norm1"] = std::make_shared(config.dim, 1e-5f); + blocks["cross_attn_norm2"] = std::make_shared(config.dim, 1e-5f); + } + } + + void init_params(ggml_context* ctx, const String2TensorStorage& weights = {}, const std::string prefix = "") override { + if (injection) { + params["cross_attn_gate"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, ggml_tensor* noisy, ggml_tensor* clean, int image_start, int image_count, ggml_tensor* text = nullptr) { + auto c = ctx->ggml_ctx; + auto norm = [&](const char* name, ggml_tensor* h) { + return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); + }; + auto modulation = std::dynamic_pointer_cast(blocks["adaLN_modulation.0"]); + auto mn = ggml_ext_chunk(c, modulation->forward(ctx, noisy), 4, 0); + auto mc = ggml_ext_chunk(c, modulation->forward(ctx, clean), 4, 0); + auto select = [&](int index, bool gate) { + auto n = gate ? ggml_tanh(c, mn[index]) : mn[index]; + auto k = gate ? ggml_tanh(c, mc[index]) : mc[index]; + auto repeat = [&](ggml_tensor* h, int64_t count) { + return ggml_repeat(c, h, ggml_new_tensor_2d(c, GGML_TYPE_F32, h->ne[0], count)); + }; + auto m = repeat(n, image_count); + if (image_start) { + m = ggml_concat(c, repeat(k, image_start), m, 1); + } + int64_t tail = x->ne[1] - image_start - image_count; + if (tail) { + m = ggml_concat(c, m, repeat(k, tail), 1); + } + return m; + }; + auto h = ZImage::modulate(ctx, norm("attention_norm1", x), select(0, false)); + h = std::dynamic_pointer_cast(blocks["attention"])->forward(ctx, h, h, pe); + x = ggml_add(c, x, ggml_mul(c, norm("attention_norm2", h), select(1, true))); + if (injection) { + h = std::dynamic_pointer_cast(blocks["cross_attention"])->forward(ctx, norm("cross_attn_norm1", x), text); + h = ggml_mul(c, norm("cross_attn_norm2", h), ggml_tanh(c, params["cross_attn_gate"])); + x = ggml_add(c, x, h); + } + h = ZImage::modulate(ctx, norm("ffn_norm1", x), select(2, false)); + h = std::dynamic_pointer_cast(blocks["feed_forward"])->forward(ctx, h); + return ggml_add(c, x, ggml_mul(c, norm("ffn_norm2", h), select(3, true))); + } + }; + + struct Fibo2Model : public GGMLBlock { + Fibo2Config config; + + Fibo2Model(const Fibo2Config& config) + : config(config) { + blocks["x_embedder"] = std::make_shared(config.channels, config.dim); + blocks["t_embedder"] = std::make_shared(1024, 256, std::min(config.dim, 256)); + blocks["perceiver"] = std::make_shared(config); + for (int i = 0; i < config.refiner_depth; ++i) { + blocks["noise_refiner." + std::to_string(i)] = std::make_shared(config, false); + } + for (int i = 0; i < config.depth; ++i) { + bool injection = std::find(config.injection_layers.begin(), config.injection_layers.end(), i) != config.injection_layers.end(); + blocks["layers." + std::to_string(i)] = std::make_shared(config, injection); + } + blocks["final_layer"] = std::make_shared(config.dim, 1, config.channels); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* t, ggml_tensor* text, ggml_tensor* pe, int gist_count, const std::vector& refs) { + auto c = ctx->ggml_ctx; + int64_t width = x->ne[0], height = x->ne[1]; + int image_count = static_cast(width * height); + // The shared FLUX.2 VAE already packs (channel, dy, dx) into 128 channels. + auto tokens = [&](ggml_tensor* h) { + h = ggml_reshape_3d(c, h, h->ne[0] * h->ne[1], config.channels, 1); + return ggml_cont(c, ggml_permute(c, h, 1, 0, 2, 3)); + }; + x = tokens(x); + for (auto ref : refs) { + x = ggml_concat(c, x, tokens(ref), 1); + } + x = std::dynamic_pointer_cast(blocks["x_embedder"])->forward(ctx, x); + auto embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); + auto noisy = embedder->forward(ctx, t); + auto clean = embedder->forward(ctx, ggml_scale(c, t, 0.f)); + auto image_pe = ggml_ext_slice(c, pe, 3, gist_count, pe->ne[3]); + for (int i = 0; i < config.refiner_depth; ++i) { + x = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)])->forward(ctx, x, image_pe, noisy, clean, 0, image_count); + sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.noise_refiner." + std::to_string(i), "x"); + } + auto bundle = [&](int i) { return ggml_ext_slice(c, text, 0, i * config.text_dim, (i + 1) * config.text_dim); }; + auto gist = std::dynamic_pointer_cast(blocks["perceiver"])->forward(ctx, bundle(0), gist_count); + x = ggml_concat(c, gist, x, 1); + int bundle_index = 1; + for (int i = 0; i < config.depth; ++i) { + auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); + auto context = block->injection ? bundle(bundle_index++) : nullptr; + x = block->forward(ctx, x, pe, noisy, clean, gist_count, image_count, context); + sd::ggml_graph_cut::mark_graph_cut(x, "fibo2.layers." + std::to_string(i), "x"); + } + x = ggml_ext_slice(c, x, 1, gist_count, gist_count + image_count); + x = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, x, noisy); + x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3)); + return ggml_reshape_4d(c, x, width, height, config.channels, 1); + } + }; + + struct Fibo2Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; + Fibo2Config config; + Fibo2Model model; + + Fibo2Runner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) + : DiffusionModelRunner(backend, prefix, weight_manager), + config(Fibo2Config::detect_from_weights(weights, prefix)), + model(config) { + model.init(params_ctx, weights, prefix); + } + + std::string get_desc() override { return "fibo2"; } + + void get_param_tensors(std::map& tensors, const std::string& prefix) override { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const DiffusionParams& p) override { + if (!p.x || !p.timesteps || !p.context || p.context->shape()[0] != config.text_dim * 6) { + LOG_ERROR("Fibo 2.0 requires six text feature bundles"); + return {}; + } + auto build = [&]() { + auto graph = new_graph_custom(65536); + auto x = make_input(*p.x); + GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels); + int count = config.gist_count(static_cast(p.context->shape()[1])); + std::vector refs; + if (p.ref_latents && p.ref_image_params.pass_to_dit) { + for (const auto& ref : *p.ref_latents) { + refs.push_back(make_input(ref)); + } + } + sd::PECacheKey pe_key("fibo2.pe"); + pe_key.add_shape(x).add(count).add(config.max_gist).add_shapes(refs).add(circular_x_enabled); + pe_key.add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + Rope::Embedding embedding; + for (int i = 0; i < count; ++i) { + embedding.ids.push_back({static_cast(i + 1), 0.f, 0.f}); + } + embedding.positions.append_tokens(count); + auto append_image = [&](ggml_tensor* image, int plane) { + auto ids = Rope::gen_flux_img_ids(static_cast(image->ne[1]), static_cast(image->ne[0]), + 1, 1, 3, plane, 0, 0, false, &embedding.positions); + embedding.ids.insert(embedding.ids.end(), ids.begin(), ids.end()); + }; + append_image(x, static_cast(config.max_gist + 1)); + for (size_t i = 0; i < refs.size(); ++i) { + append_image(refs[i], static_cast(config.max_gist + 2 + i)); + } + embedding.values = Rope::embed_nd(embedding.ids, 1, 256.f, {32, 48, 48}, embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + const int64_t tokens = static_cast(values.size() / 256); + return sd::Tensor({2, 2, 64, tokens}, std::move(values)); + })); + auto ctx = get_context(); + auto out = model.forward(&ctx, x, make_input(*p.timesteps), make_input(*p.context), pe, count, refs); + ggml_build_forward_expand(graph, out); + return graph; + }; + return restore_trailing_singleton_dims(GGMLRunner::compute(build, n_threads, false), p.x->dim()); + } + }; +} + +#endif // __SD_MODEL_DIFFUSION_FIBO2_HPP__ diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 825dabe13d..40edca0ed9 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -6,6 +6,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "core/util.h" #include "model/adapter/pulid.hpp" @@ -76,7 +77,10 @@ namespace Flux { config.in_channels = 128; } else if (version == VERSION_FLEX_2) { config.in_channels = 196; + } else if (version == VERSION_CHROMA) { + config.is_chroma = true; } else if (version == VERSION_CHROMA_RADIANCE) { + config.is_chroma = true; config.in_channels = 3; config.patch_size = 16; } else if (version == VERSION_OVIS_IMAGE) { @@ -137,9 +141,6 @@ namespace Flux { actual_radiance_patch_size = tensor_storage.ne[0]; LOG_VERBOSE("actual radiance patch size: %" PRId64, actual_radiance_patch_size); } - if (name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { - config.is_chroma = true; - } size_t db = name.find("double_blocks."); if (db != std::string::npos) { std::string block_name = name.substr(db); @@ -357,9 +358,8 @@ namespace Flux { auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); auto gate = gate_proj->forward(ctx, x); - gate = ggml_silu_inplace(ctx->ggml_ctx, gate); x = up_proj->forward(ctx, x); - x = ggml_mul(ctx->ggml_ctx, x, gate); + x = ggml_swiglu_split(ctx->ggml_ctx, gate, x); x = down_proj->forward(ctx, x); return x; } @@ -413,7 +413,7 @@ namespace Flux { } }; - __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* shift, ggml_tensor* scale, @@ -422,12 +422,10 @@ namespace Flux { // scale: [N, C] // shift: [N, C] if (!skip_reshape) { - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] - shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C] + scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] + shift = ggml_reshape_3d(ctx->ggml_ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C] } - x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); - x = ggml_add(ctx, x, shift); - return x; + return ggml_ext_modulate(ctx, x, shift, scale); } struct DoubleStreamBlock : public GGMLBlock { @@ -540,7 +538,7 @@ namespace Flux { // prepare image for attention auto img_modulated = img_norm1->forward(ctx, img); - img_modulated = Flux::modulate(ctx->ggml_ctx, img_modulated, img_mod1.shift, img_mod1.scale); + img_modulated = Flux::modulate(ctx, img_modulated, img_mod1.shift, img_mod1.scale); auto img_qkv = img_attn->pre_attention(ctx, img_modulated); // q,k,v: [N, n_img_token, n_head, d_head] auto img_q = img_qkv[0]; auto img_k = img_qkv[1]; @@ -548,7 +546,7 @@ namespace Flux { // prepare txt for attention auto txt_modulated = txt_norm1->forward(ctx, txt); - txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_modulated, txt_mod1.shift, txt_mod1.scale); + txt_modulated = Flux::modulate(ctx, txt_modulated, txt_mod1.shift, txt_mod1.scale); auto txt_qkv = txt_attn->pre_attention(ctx, txt_modulated); // q,k,v: [N, n_txt_token, n_head, d_head] auto txt_q = txt_qkv[0]; auto txt_k = txt_qkv[1]; @@ -580,14 +578,14 @@ namespace Flux { // calculate the img bloks img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_attn->post_attention(ctx, img_attn_out), img_mod1.gate)); - auto img_mlp_out = img_mlp->forward(ctx, Flux::modulate(ctx->ggml_ctx, img_norm2->forward(ctx, img), img_mod2.shift, img_mod2.scale)); + auto img_mlp_out = img_mlp->forward(ctx, Flux::modulate(ctx, img_norm2->forward(ctx, img), img_mod2.shift, img_mod2.scale)); img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp_out, img_mod2.gate)); // calculate the txt bloks txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn->post_attention(ctx, txt_attn_out), txt_mod1.gate)); - auto txt_mlp_out = txt_mlp->forward(ctx, Flux::modulate(ctx->ggml_ctx, txt_norm2->forward(ctx, txt), txt_mod2.shift, txt_mod2.scale)); + auto txt_mlp_out = txt_mlp->forward(ctx, Flux::modulate(ctx, txt_norm2->forward(ctx, txt), txt_mod2.shift, txt_mod2.scale)); txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp_out, txt_mod2.gate)); return {img, txt}; @@ -670,7 +668,7 @@ namespace Flux { } } - auto x_mod = Flux::modulate(ctx->ggml_ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale); + auto x_mod = Flux::modulate(ctx, pre_norm->forward(ctx, x), mod.shift, mod.scale); auto qkv_mlp = linear1->forward(ctx, x_mod); // [N, n_token, hidden_size * 3 + mlp_hidden_dim*mlp_mult_factor] int64_t head_dim = hidden_size / num_heads; @@ -752,7 +750,7 @@ namespace Flux { scale = m_vec[1]; } - x = Flux::modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale); + x = Flux::modulate(ctx, norm_final->forward(ctx, x), shift, scale); x = linear->forward(ctx, x); return x; @@ -849,12 +847,10 @@ namespace Flux { auto res_x = x; x = norm->forward(ctx, x); // [batch_size, n_token, hidden_size_x] - auto x1 = ggml_mul_mat(ctx->ggml_ctx, fc1_gate, x); // [batch_size, n_token, hidden_size_x*mlp_ratio] - x1 = ggml_silu_inplace(ctx->ggml_ctx, x1); - + auto x1 = ggml_mul_mat(ctx->ggml_ctx, fc1_gate, x); // [batch_size, n_token, hidden_size_x*mlp_ratio] auto x2 = ggml_mul_mat(ctx->ggml_ctx, fc1_value, x); // [batch_size, n_token, hidden_size_x*mlp_ratio] - x = ggml_mul_inplace(ctx->ggml_ctx, x1, x2); // [batch_size, n_token, hidden_size_x*mlp_ratio] + x = ggml_swiglu_split(ctx->ggml_ctx, x1, x2); // [batch_size, n_token, hidden_size_x*mlp_ratio] x = ggml_mul_mat(ctx->ggml_ctx, fc2, x); // [batch_size, n_token, hidden_size_x] @@ -1398,10 +1394,11 @@ namespace Flux { }; struct FluxRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: FluxConfig config; Flux flux; - std::vector pe_vec; std::vector mod_index_arange_vec; std::vector dct_vec; sd::Tensor guidance_tensor; @@ -1548,25 +1545,26 @@ namespace Flux { } else if (version == VERSION_OVIS_IMAGE) { txt_arange_dims = {1, 2}; } - pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - txt_arange_dims, - ref_latents, - ref_index_mode, - config.ref_index_scale, - config.theta, - config.axes_dim, - sd_version_is_longcat(version))); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("flux.pe"); + pe_key.add(version).add_shape(x).add_shape(context).add(config.patch_size).add(txt_arange_dims); + pe_key.add_shapes(ref_latents).add(ref_index_mode).add(config.ref_index_scale).add(config.theta).add(config.axes_dim); + pe_key.add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_flux_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + txt_arange_dims, + ref_latents, + ref_index_mode, + config.ref_index_scale, + config.theta, + config.axes_dim, + sd_version_is_longcat(version))); + int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); if (version == VERSION_CHROMA_RADIANCE) { int patch_size = config.patch_size; diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp index 952587b013..fac1b1a567 100644 --- a/src/model/diffusion/flux3.hpp +++ b/src/model/diffusion/flux3.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_FLUX3_HPP__ #define __SD_MODEL_DIFFUSION_FLUX3_HPP__ +#include "core/pe_cache.h" #include "model/diffusion/flux.hpp" namespace Flux3 { @@ -82,7 +83,7 @@ namespace Flux3 { int64_t offset = 0; for (size_t i = 0; i < lengths.size(); ++i) { auto part = token_slice(gc, x, offset, lengths[i]); - part = Flux::modulate(gc, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale); + part = Flux::modulate(ctx, ggml_norm(gc, part, 1e-6f), mods[i].shift, mods[i].scale); x_mod = x_mod ? ggml_concat(gc, x_mod, part, 1) : part; offset += lengths[i]; } @@ -170,12 +171,17 @@ namespace Flux3 { const std::string base = mode == "txt" ? "txt_mode_blocks." : "content_mode_blocks." + mode + "."; for (int i = 0; i < config.depth; ++i) { h = std::dynamic_pointer_cast(blocks[base + std::to_string(i)])->forward(ctx, h, mode_pe, {early}, {lengths.back()}); + if (i + 1 < config.depth) { + sd::ggml_graph_cut::mark_graph_cut(h, "flux3." + base + std::to_string(i), "h"); + } } joint = joint ? ggml_concat(gc, joint, h, 1) : h; + sd::ggml_graph_cut::mark_graph_cut(joint, "flux3." + base + "joint", "joint"); offset += lengths.back(); } for (int i = 0; i < config.single_depth; ++i) { joint = std::dynamic_pointer_cast(blocks["single_blocks." + std::to_string(i)])->forward(ctx, joint, pe, single_mods, lengths); + sd::ggml_graph_cut::mark_graph_cut(joint, "flux3.single_blocks." + std::to_string(i), "joint"); } auto out = std::dynamic_pointer_cast(blocks["final_layer.image"])->forward(ctx, token_slice(gc, joint, 0, lengths[0]), image_vec); return DiT::unpatchify(gc, out, height, width, 1, 1); @@ -183,9 +189,9 @@ namespace Flux3 { }; struct Flux3Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; Flux3Config config; Model model; - std::vector pe_values; sd::Tensor zero = sd::Tensor::zeros({1}); Flux3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) @@ -213,37 +219,44 @@ namespace Flux3 { auto x = make_input(*params.x); auto text = make_input(*params.context); GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1); - Rope::Embedding embedding; - auto append_image = [&](int64_t w, int64_t h, float index) { - embedding.positions.append_image(static_cast(h), static_cast(w), 1, 0.5f, 0.5f); - for (int64_t y = 0; y < h; ++y) { - for (int64_t x = 0; x < w; ++x) { - embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index}); - } - } - }; - append_image(x->ne[0], x->ne[1], 0.f); + std::vector refs; ggml_tensor* references = nullptr; if (params.ref_latents && params.ref_image_params.pass_to_dit) { - int index = 0; for (const auto& ref : *params.ref_latents) { auto r = make_input(ref); - append_image(r->ne[0], r->ne[1], -float(++index)); + refs.push_back(r); r = DiT::patchify(compute_ctx, r, 1, 1); references = references ? ggml_concat(compute_ctx, references, r, 1) : r; } } - for (int64_t i = 0; i < text->ne[1]; ++i) { - embedding.ids.push_back({0.f, 0.f, 0.f, static_cast(i)}); - } - embedding.positions.append_tokens(text->ne[1]); - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies); - pe_values = finish_rope_pe(std::move(embedding)); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, 64, pe_values.size() / 256); - set_backend_tensor_data(pe, pe_values.data()); + sd::PECacheKey pe_key("flux3.pe"); + pe_key.add_shape(x).add_shape(text).add_shapes(refs).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + Rope::Embedding embedding; + auto append_image = [&](int64_t w, int64_t h, float index) { + embedding.positions.append_image(static_cast(h), static_cast(w), 1, 0.5f, 0.5f); + for (int64_t y = 0; y < h; ++y) { + for (int64_t x = 0; x < w; ++x) { + embedding.ids.push_back({0.f, y * 0.5f, x * 0.5f, index}); + } + } + }; + append_image(x->ne[0], x->ne[1], 0.f); + for (size_t i = 0; i < refs.size(); ++i) { + append_image(refs[i]->ne[0], refs[i]->ne[1], -float(i + 1)); + } + for (int64_t i = 0; i < text->ne[1]; ++i) { + embedding.ids.push_back({0.f, 0.f, 0.f, static_cast(i)}); + } + embedding.positions.append_tokens(text->ne[1]); + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, {32, 32, 32, 32}, embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + const int64_t tokens = static_cast(values.size() / 256); + return sd::Tensor({2, 2, 64, tokens}, std::move(values)); + })); auto runner_ctx = get_context(); auto out = model.forward(&runner_ctx, x, text, references, make_input(*params.timesteps), make_input(zero), - config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe); + config.flow_map ? make_input(*extra->target_timestep) : nullptr, pe); ggml_build_forward_expand(gf, out); return gf; }; diff --git a/src/model/diffusion/hidream_o1.hpp b/src/model/diffusion/hidream_o1.hpp index 2befb723a3..90bfb51b34 100644 --- a/src/model/diffusion/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -56,15 +56,6 @@ namespace HiDreamO1 { } }; - static inline std::string repeat_special_token(const std::string& token, int64_t count) { - std::string out; - out.reserve(static_cast(count) * token.size()); - for (int64_t i = 0; i < count; ++i) { - out += token; - } - return out; - } - static inline std::pair calculate_dimensions(int max_size, double ratio) { int width = static_cast(std::sqrt(max_size * max_size * ratio)); int height = static_cast(width / ratio); @@ -247,10 +238,10 @@ namespace HiDreamO1 { blocks["language_model"] = std::make_shared(this->config.llm); blocks["t_embedder1"] = std::make_shared(this->config.llm.hidden_size); blocks["x_embedder"] = std::make_shared(this->config.patch_size * this->config.patch_size * 3, - this->config.llm.hidden_size / 4, - this->config.llm.hidden_size); + this->config.llm.hidden_size / 4, + this->config.llm.hidden_size); blocks["final_layer2"] = std::make_shared(this->config.llm.hidden_size, - this->config.patch_size * this->config.patch_size * 3); + this->config.patch_size * this->config.patch_size * 3); } std::shared_ptr text_model() { @@ -405,7 +396,7 @@ namespace HiDreamO1 { txt = LLM::splice_image_embeds(&runner_ctx, txt, image_embeds); auto t_emb = t_embedder1->forward(&runner_ctx, timestep); - int64_t txt_seq_len = input_ids->ne[0]; + int64_t txt_seq_len = txt->ne[1]; if (txt_seq_len > 1) { auto prefix = ggml_ext_slice(compute_ctx, txt, 1, 0, txt_seq_len - 1); txt = ggml_concat(compute_ctx, prefix, ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1), 1); @@ -616,8 +607,7 @@ namespace HiDreamO1 { sd::Tensor vlm_image = sd::ops::interpolate( resized_ref, {dims.first, dims.second, resized_ref.shape()[2], resized_ref.shape()[3]}); - vlm_image = vlm_image * 2.0f - 1.0f; - int64_t image_tokens = static_cast(dims.first / PATCH_SIZE) * static_cast(dims.second / PATCH_SIZE); + vlm_image = vlm_image * 2.0f - 1.0f; auto patch_img = resized_ref * 2.0f - 1.0f; result.c_ref_images.push_back(std::move(patch_img)); @@ -626,8 +616,11 @@ namespace HiDreamO1 { return {}; } int64_t prompt_start = static_cast(prefix_tokens.size()); + for (const auto& grid : image_grids) { + prompt_start += static_cast(grid[1]) * grid[2] - 1; + } prompt += "<|vision_start|>"; - prompt += repeat_special_token("<|image_pad|>", image_tokens); + prompt += "<|image_pad|>"; prompt += "<|vision_end|>"; vlm_images.emplace_back(static_cast(prompt_start), std::move(vlm_image)); image_grids.push_back({1, dims.second / PATCH_SIZE, dims.first / PATCH_SIZE}); @@ -642,6 +635,13 @@ namespace HiDreamO1 { } std::vector input_ids_pad = input_ids; + for (size_t i = 0; i < vlm_images.size(); ++i) { + const auto& grid = image_grids[i]; + const int64_t count = static_cast(grid[1]) * grid[2]; + input_ids_pad.insert(input_ids_pad.begin() + vlm_images[i].first, count - 1, IMAGE_TOKEN_ID); + } + + int txt_seq_len = static_cast(input_ids_pad.size()); input_ids_pad.push_back(VISION_START_TOKEN_ID); input_ids_pad.insert(input_ids_pad.end(), target_image_len - 1, IMAGE_TOKEN_ID); image_grids.push_back({1, static_cast(height / PATCH_SIZE), static_cast(width / PATCH_SIZE)}); @@ -656,8 +656,7 @@ namespace HiDreamO1 { } std::vector token_types(input_ids_pad.size(), 0); - int txt_seq_len = static_cast(input_ids.size()); - int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; + int bgn = txt_seq_len - TIMESTEP_TOKEN_NUM; for (int i = bgn; i < static_cast(token_types.size()); ++i) { token_types[i] = 1; } diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index c106b9cd59..5fc29b6b0a 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/diffusion/flux.hpp" @@ -302,9 +303,9 @@ namespace Hunyuan { int64_t head_dim = this->config.hidden_size / this->config.num_heads; blocks["txt_in"] = std::make_shared(this->config.context_in_dim, this->config.num_heads, head_dim, 2); blocks["img_in"] = std::make_shared(static_cast(224) /*Not used*/, - this->config.patch_size, - this->config.in_channels, - this->config.hidden_size); + this->config.patch_size, + this->config.in_channels, + this->config.hidden_size); blocks["time_in"] = std::make_shared(256, this->config.hidden_size); if (this->config.vec_in_dim > 0) { blocks["vector_in"] = std::make_shared(this->config.vec_in_dim, this->config.hidden_size); @@ -540,10 +541,11 @@ namespace Hunyuan { }; struct HunyuanVideoRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: HunyuanVideoConfig config; HunyuanVideoModel hunyuan_video; - std::vector pe_vec; SDVersion version; HunyuanVideoRunner(ggml_backend_t backend, @@ -605,23 +607,24 @@ namespace Hunyuan { if (vision != nullptr) { text_len += static_cast(vision->ne[1]); } - pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(x->ne[0]), - std::get<0>(config.patch_size), - std::get<1>(config.patch_size), - std::get<2>(config.patch_size), - 1, - text_len, - config.theta, - config.axes_dim); - int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe, true, "pe"); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("hunyuan.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(std::get<0>(config.patch_size)); + pe_key.add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size)).add(text_len); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = Rope::gen_hunyuan_video_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + text_len, + config.theta, + config.axes_dim); + int64_t pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 3254571db2..270344dc97 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -8,6 +8,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" @@ -102,9 +103,9 @@ namespace Ideogram4 { return ggml_reshape_3d(ctx, x, hidden_size * LLM_HIDDEN_STATE_LAYERS, token_count, batch_count); } - __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* scale) { - scale = to_token_modulation(ctx, scale); - return ggml_add(ctx, x, ggml_mul(ctx, x, scale)); + __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) { + scale = to_token_modulation(ctx->ggml_ctx, scale); + return ggml_ext_modulate(ctx, x, nullptr, scale); } __STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) { @@ -243,9 +244,9 @@ namespace Ideogram4 { auto w2 = std::dynamic_pointer_cast(blocks["w2"]); auto w3 = std::dynamic_pointer_cast(blocks["w3"]); - auto x1 = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x)); + auto x1 = w1->forward(ctx, x); auto x3 = w3->forward(ctx, x); - x = ggml_mul(ctx->ggml_ctx, x1, x3); + x = ggml_swiglu_split(ctx->ggml_ctx, x1, x3); x = w2->forward(ctx, x); return x; } @@ -284,13 +285,13 @@ namespace Ideogram4 { auto gate_mlp = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[3])); auto attn_out = attention_norm1->forward(ctx, x); - attn_out = modulate(ctx->ggml_ctx, attn_out, scale_msa); + attn_out = modulate(ctx, attn_out, scale_msa); attn_out = attention->forward(ctx, attn_out, pe, mask); attn_out = attention_norm2->forward(ctx, attn_out); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa)); auto ffn_out = ffn_norm1->forward(ctx, x); - ffn_out = modulate(ctx->ggml_ctx, ffn_out, scale_mlp); + ffn_out = modulate(ctx, ffn_out, scale_mlp); ffn_out = feed_forward->forward(ctx, ffn_out); ffn_out = ffn_norm2->forward(ctx, ffn_out); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, ffn_out, gate_mlp)); @@ -336,7 +337,7 @@ namespace Ideogram4 { auto scale = adaln_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); x = norm_final->forward(ctx, x); - x = modulate(ctx->ggml_ctx, x, scale); + x = modulate(ctx, x, scale); x = linear->forward(ctx, x); return x; } @@ -422,6 +423,8 @@ namespace Ideogram4 { }; class Ideogram4Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; + protected: bool should_use_uncond_model(const DiffusionParams& diffusion_params) const { return has_uncond_model && @@ -436,7 +439,6 @@ namespace Ideogram4 { Ideogram4Transformer uncond_model; bool has_uncond_model = false; std::string uncond_prefix; - std::vector pe_vec; std::vector image_indicator_vec; Ideogram4Runner(ggml_backend_t backend, @@ -496,15 +498,19 @@ namespace Ideogram4 { int64_t head_dim = config.emb_dim / config.num_heads; auto runner_ctx = get_context(); - pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), + sd::PECacheKey pe_key("ideogram4.pe"); + pe_key.add(grid_h).add(grid_w).add(x->ne[3]).add(context_len).add(head_dim).add(pos_len).add(config.rope_theta); + pe_key.add(config.mrope_section).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(gen_ideogram4_pe(static_cast(grid_h), static_cast(grid_w), static_cast(x->ne[3]), static_cast(context_len), static_cast(head_dim), static_cast(config.rope_theta), config.mrope_section)); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, head_dim / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + return sd::Tensor({2, 2, head_dim / 2, pos_len}, std::move(pe_vec)); + })); image_indicator_vec.assign(static_cast(pos_len), 1); for (int64_t i = 0; i < context_len; ++i) { diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp index b90fbe9d44..3681c65ab2 100644 --- a/src/model/diffusion/iris.hpp +++ b/src/model/diffusion/iris.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/diffusion/pid.hpp" namespace Iris { @@ -382,11 +383,9 @@ namespace Iris { } struct IrisRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; IrisConfig config; IrisModel model; - std::vector joint_pe_data; - std::vector pixel_pe_data; - std::vector pixel_pos_data; std::vector text_mask_data; IrisRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) @@ -418,24 +417,35 @@ namespace Iris { GGML_ASSERT(h % config.patch_size == 0 && w % config.patch_size == 0); GGML_ASSERT(text->ne[0] == config.text_dim * config.text_layers && text->ne[1] == nt); GGML_ASSERT(inputs.y->numel() == nt && x->ne[3] == 1); - joint_pe_data = Pid::make_rope_1d(nt, static_cast(config.head_dim), 10000.f); - auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim); - joint_pe_data.insert(joint_pe_data.end(), img.begin(), img.end()); - pixel_pe_data = image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads); - pixel_pos_data = Pid::make_pixel_abs_pos(h, w, static_cast(config.pixel_dim)); + sd::PECacheKey joint_key("iris.joint_pe"); + joint_key.add(h).add(w).add(nt).add(config.patch_size).add(config.head_dim); + auto pe = make_input(pe_cache_.get_or_create(joint_key, [&]() { + auto values = Pid::make_rope_1d(nt, static_cast(config.head_dim), 10000.f); + auto img = image_rope(h / config.patch_size, w / config.patch_size, config.head_dim); + values.insert(values.end(), img.begin(), img.end()); + const int64_t tokens = nt + static_cast(img.size() / (2 * config.head_dim)); + return sd::Tensor({2, 2, config.head_dim / 2, tokens}, std::move(values)); + })); + sd::PECacheKey pixel_key("iris.pixel_pe"); + pixel_key.add(h).add(w).add(config.patch_size).add(config.pixel_attn_dim).add(config.pixel_heads); + auto pp = make_input(pe_cache_.get_or_create(pixel_key, [&]() { + return sd::Tensor({2, 2, config.pixel_attn_dim / config.pixel_heads / 2, + (h / config.patch_size) * (w / config.patch_size)}, + image_rope(h / config.patch_size, w / config.patch_size, config.pixel_attn_dim / config.pixel_heads)); + })); + sd::PECacheKey pos_key("iris.pixel_pos"); + pos_key.add(h).add(w).add(config.pixel_dim); + auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() { + return sd::Tensor({config.pixel_dim, h * w, 1}, + Pid::make_pixel_abs_pos(h, w, static_cast(config.pixel_dim))); + })); text_mask_data.resize(nt * nt); for (int q = 0; q < nt; ++q) { for (int k = 0; k < nt; ++k) { text_mask_data[q * nt + k] = inputs.y->values()[k] != 0.f || q == k ? 0.f : -INFINITY; } } - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, nt + img.size() / (2 * config.head_dim)); - auto pp = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.pixel_attn_dim / config.pixel_heads / 2, (h / config.patch_size) * (w / config.patch_size)); - auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.pixel_dim, h * w, 1); auto mask = ggml_new_tensor_2d(compute_ctx, GGML_TYPE_F32, nt, nt); - set_backend_tensor_data(pe, joint_pe_data.data()); - set_backend_tensor_data(pp, pixel_pe_data.data()); - set_backend_tensor_data(pos, pixel_pos_data.data()); set_backend_tensor_data(mask, text_mask_data.data()); auto ctx = get_context(); auto out = model.forward(&ctx, x, t, text, mask, pe, pp, pos); diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index e58df41081..fba817aa67 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -11,6 +11,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_graph_cut.h" @@ -73,9 +74,9 @@ namespace Krea2 { int64_t dim_head = head_dim(); int64_t unit = dim_head / 16; axes_dim = { - static_cast(dim_head - 12 * unit), - static_cast(6 * unit), - static_cast(6 * unit), + static_cast(dim_head - 12 * unit), + static_cast(6 * unit), + static_cast(6 * unit), }; axes_dim_sum = axes_dim[0] + axes_dim[1] + axes_dim[2]; } @@ -207,9 +208,9 @@ namespace Krea2 { auto up = std::dynamic_pointer_cast(blocks["up"]); auto down = std::dynamic_pointer_cast(blocks["down"]); - auto gated = ggml_silu(ctx->ggml_ctx, gate->forward(ctx, x)); - auto up_x = up->forward(ctx, x); - x = ggml_mul(ctx->ggml_ctx, gated, up_x); + auto gate_x = gate->forward(ctx, x); + auto up_x = up->forward(ctx, x); + x = ggml_swiglu_split(ctx->ggml_ctx, gate_x, up_x); return down->forward(ctx, x); } }; @@ -387,10 +388,10 @@ namespace Krea2 { this->config.mlp_multiplier, this->config.norm_eps); blocks["refiner_blocks." + std::to_string(i)] = std::make_shared(this->config.text_dim, - this->config.text_heads, - this->config.text_kv_heads, - this->config.mlp_multiplier, - this->config.norm_eps); + this->config.text_heads, + this->config.text_kv_heads, + this->config.mlp_multiplier, + this->config.norm_eps); } blocks["projector"] = std::make_shared(this->config.text_layers, 1, false); } @@ -464,8 +465,8 @@ namespace Krea2 { auto pre_x_main = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_main, B, nb1, nb2, 0); auto pre_x_refs = ggml_view_3d(ctx->ggml_ctx, pre_x, D, len_refs, B, nb1, nb2, len_main * nb1); - auto attn_in_main = Flux::modulate(ctx->ggml_ctx, pre_x_main, mods_main[1], mods_main[0], true); - auto attn_in_refs = Flux::modulate(ctx->ggml_ctx, pre_x_refs, mods_refs[1], mods_refs[0], true); + auto attn_in_main = Flux::modulate(ctx, pre_x_main, mods_main[1], mods_main[0], true); + auto attn_in_refs = Flux::modulate(ctx, pre_x_refs, mods_refs[1], mods_refs[0], true); auto attn_input = ggml_concat(ctx->ggml_ctx, attn_in_main, attn_in_refs, 1); @@ -486,8 +487,8 @@ namespace Krea2 { auto post_x_main = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_main, B, post_x->nb[1], post_x->nb[2], 0); auto post_x_refs = ggml_view_3d(ctx->ggml_ctx, post_x, D, len_refs, B, post_x->nb[1], post_x->nb[2], len_main * post_x->nb[1]); - auto mlp_in_main = Flux::modulate(ctx->ggml_ctx, post_x_main, mods_main[4], mods_main[3], true); - auto mlp_in_refs = Flux::modulate(ctx->ggml_ctx, post_x_refs, mods_refs[4], mods_refs[3], true); + auto mlp_in_main = Flux::modulate(ctx, post_x_main, mods_main[4], mods_main[3], true); + auto mlp_in_refs = Flux::modulate(ctx, post_x_refs, mods_refs[4], mods_refs[3], true); auto mlp_input = ggml_concat(ctx->ggml_ctx, mlp_in_main, mlp_in_refs, 1); auto mlp_out = mlp->forward(ctx, mlp_input); @@ -502,7 +503,7 @@ namespace Krea2 { x = ggml_add(ctx->ggml_ctx, x, mlp_res); } else { auto mods = mod->forward(ctx, vec); - auto attn_input = Flux::modulate(ctx->ggml_ctx, + auto attn_input = Flux::modulate(ctx, prenorm->forward(ctx, x), mods[1], mods[0], @@ -510,7 +511,7 @@ namespace Krea2 { auto attn_out = attn->forward(ctx, attn_input, pe); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, mods[2])); - auto mlp_input = Flux::modulate(ctx->ggml_ctx, + auto mlp_input = Flux::modulate(ctx, postnorm->forward(ctx, x), mods[4], mods[3], @@ -594,7 +595,7 @@ namespace Krea2 { auto modulation = std::dynamic_pointer_cast(blocks["modulation"]); auto mods = modulation->forward(ctx, vec); - x = Flux::modulate(ctx->ggml_ctx, + x = Flux::modulate(ctx, norm->forward(ctx, x), mods[1], mods[0], @@ -613,8 +614,8 @@ namespace Krea2 { explicit Krea2Model(Krea2Config config) : config(std::move(config)) { blocks["first"] = std::make_shared(this->config.patch_size * this->config.patch_size * this->config.in_channels, - this->config.features, - true); + this->config.features, + true); blocks["tmlp"] = std::make_shared(this->config); blocks["txtfusion"] = std::make_shared(this->config); blocks["txtmlp"] = std::make_shared(this->config); @@ -714,9 +715,9 @@ namespace Krea2 { } struct Krea2Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; Krea2Config config; Krea2Model model; - std::vector pe_vec; Krea2Runner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -754,18 +755,23 @@ namespace Krea2 { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = finish_rope_pe(gen_krea2_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim, - ref_latents, - ref_image_params.ref_index_mode)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("krea2.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents); + pe_key.add(ref_image_params.ref_index_mode); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(gen_krea2_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim, + ref_latents, + ref_image_params.ref_index_mode)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pe, ref_latents, ref_image_params.force_ref_timestep_zero); diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 8fb1512b41..bcfefd97bf 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -4,6 +4,7 @@ #include #include #include +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -103,9 +104,9 @@ namespace Lens { auto w2 = std::dynamic_pointer_cast(blocks["w2"]); auto w3 = std::dynamic_pointer_cast(blocks["w3"]); - auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x)); + auto gate = w1->forward(ctx, x); auto up = w3->forward(ctx, x); - x = ggml_mul(ctx->ggml_ctx, gate, up); + x = ggml_swiglu_split(ctx->ggml_ctx, gate, up); return w2->forward(ctx, x); } }; @@ -236,9 +237,9 @@ namespace Lens { auto txt_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0); auto img_normed = img_norm1->forward(ctx, img); - auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_vec[0], img_mod_vec[1]); + auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_vec[0], img_mod_vec[1]); auto txt_normed = txt_norm1->forward(ctx, txt); - auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]); + auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_vec[0], txt_mod_vec[1]); auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe); @@ -246,9 +247,9 @@ namespace Lens { txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_mod_vec[2])); auto img_normed2 = img_norm2->forward(ctx, img); - auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]); + auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_vec[3], img_mod_vec[4]); auto txt_normed2 = txt_norm2->forward(ctx, txt); - auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]); + auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_vec[3], txt_mod_vec[4]); img = ggml_add(ctx->ggml_ctx, img, ggml_mul(ctx->ggml_ctx, img_mlp->forward(ctx, img_modulated2), img_mod_vec[5])); txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_mlp->forward(ctx, txt_modulated2), txt_mod_vec[5])); @@ -271,7 +272,7 @@ namespace Lens { auto scale = mods[0]; auto shift = mods[1]; x = ggml_norm(ctx->ggml_ctx, x, eps); - return Flux::modulate(ctx->ggml_ctx, x, shift, scale); + return Flux::modulate(ctx, x, shift, scale); } }; @@ -352,9 +353,9 @@ namespace Lens { }; struct LensRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; LensConfig config; LensModel lens; - std::vector pe_vec; LensRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -384,15 +385,19 @@ namespace Lens { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - static_cast(x->ne[3]), - static_cast(context->ne[1]), - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("lens.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_lens_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + static_cast(x->ne[3]), + static_cast(context->ne[1]), + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = lens.forward(&runner_ctx, x, timesteps, context, pe); diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index 1a93b458eb..d15d754080 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -9,6 +9,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/util.h" #include "model/common/rope.hpp" @@ -230,9 +231,8 @@ namespace LingBotVideo { auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); auto gate = gate_proj->forward(ctx, x); - gate = ggml_silu_inplace(ctx->ggml_ctx, gate); auto up = up_proj->forward(ctx, x); - x = ggml_mul(ctx->ggml_ctx, gate, up); + x = ggml_swiglu_split(ctx->ggml_ctx, gate, up); x = down_proj->forward(ctx, x); return x; } @@ -405,9 +405,8 @@ namespace LingBotVideo { x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size, 1, n_token_total); auto gate = expert_linear(ctx, "experts.w1", x, selected_experts); - gate = ggml_silu_inplace(ctx->ggml_ctx, gate); auto up = expert_linear(ctx, "experts.w3", x, selected_experts); - auto act = ggml_mul(ctx->ggml_ctx, gate, up); + auto act = ggml_swiglu_split(ctx->ggml_ctx, gate, up); auto out = expert_linear(ctx, "experts.w2", act, selected_experts); out = ggml_mul(ctx->ggml_ctx, out, weights); @@ -442,10 +441,10 @@ namespace LingBotVideo { : hidden_size(config.hidden_size) { blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["attn"] = std::make_shared(config.hidden_size, - config.num_attention_heads, - config.qkv_bias, - config.out_bias, - config.norm_eps); + config.num_attention_heads, + config.qkv_bias, + config.out_bias, + config.norm_eps); blocks["norm_post_attn"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); if (sparse) { @@ -486,12 +485,12 @@ namespace LingBotVideo { auto scale_mlp = mods[4]; auto gate_mlp = ggml_tanh(ctx->ggml_ctx, mods[5]); - auto attn_in = Flux::modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true); + auto attn_in = Flux::modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa, true); auto attn_out = attn->forward(ctx, attn_in, pe, attention_mask); attn_out = norm_post_attn->forward(ctx, attn_out); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_msa, attn_out)); - auto ffn_in = Flux::modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true); + auto ffn_in = Flux::modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp, true); auto ffn_out = ffn->forward(ctx, ffn_in); ffn_out = norm_post_ffn->forward(ctx, ffn_out); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, gate_mlp, ffn_out)); @@ -509,14 +508,14 @@ namespace LingBotVideo { int64_t patch_dim = config.patch_t * config.patch_h * config.patch_w * config.in_channels; blocks["patch_embedder"] = std::make_shared(patch_dim, config.hidden_size, config.patch_embed_bias); blocks["time_embedder"] = std::make_shared(config.freq_dim, - config.hidden_size, - config.hidden_size, - 0, - config.timestep_mlp_bias); + config.hidden_size, + config.hidden_size, + 0, + config.timestep_mlp_bias); blocks["time_modulation.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size, true); blocks["text_embedder"] = std::make_shared(config.text_dim, - config.hidden_size, - config.norm_eps); + config.hidden_size, + config.norm_eps); for (int i = 0; i < config.depth; i++) { bool sparse = config.sparse_layers.find(i) != config.sparse_layers.end(); blocks["blocks." + std::to_string(i)] = std::make_shared(config, sparse); @@ -524,8 +523,8 @@ namespace LingBotVideo { blocks["norm_out"] = std::make_shared(config.hidden_size, config.norm_eps, false); blocks["norm_out_modulation.1"] = std::make_shared(config.hidden_size, 2 * config.hidden_size, true); blocks["proj_out"] = std::make_shared(config.hidden_size, - config.patch_t * config.patch_h * config.patch_w * config.out_channels, - true); + config.patch_t * config.patch_h * config.patch_w * config.out_channels, + true); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -588,7 +587,7 @@ namespace LingBotVideo { 2, 0); hidden_states = norm_out->forward(ctx, hidden_states); - hidden_states = Flux::modulate(ctx->ggml_ctx, hidden_states, final_mods[0], final_mods[1], true); + hidden_states = Flux::modulate(ctx, hidden_states, final_mods[0], final_mods[1], true); hidden_states = proj_out->forward(ctx, hidden_states); hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img); @@ -606,9 +605,9 @@ namespace LingBotVideo { }; struct LingBotVideoRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; LingBotVideoConfig config; LingBotVideoModel lingbot_video; - std::vector pe_vec; LingBotVideoRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -647,19 +646,23 @@ namespace LingBotVideo { int64_t T = x->ne[2]; int64_t N = x_tensor.shape()[4]; GGML_ASSERT(N == 1); - pe_vec = Rope::gen_lingbot_video_pe(static_cast(T), - static_cast(H), - static_cast(W), - config.patch_t, - config.patch_h, - config.patch_w, - static_cast(N), - static_cast(context->ne[1]), - config.rope_theta, - config.axes_dim); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("lingbot_video.pe"); + pe_key.add(T).add(H).add(W).add(N).add_shape(context).add(config.patch_t).add(config.patch_h).add(config.patch_w); + pe_key.add(config.rope_theta).add(config.axes_dim).add(config.axes_dim_sum); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = Rope::gen_lingbot_video_pe(static_cast(T), + static_cast(H), + static_cast(W), + config.patch_t, + config.patch_h, + config.patch_w, + static_cast(N), + static_cast(context->ne[1]), + config.rope_theta, + config.axes_dim); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = lingbot_video.forward(&runner_ctx, x, timesteps, context, pe); diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index e69c3ecd39..4a9ea1e5ab 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -379,10 +380,11 @@ namespace LLaDAImage { }; struct LLaDAImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: LLaDAImageConfig config; LLaDAImageModel llada_image; - std::vector pe_vec; LLaDAImageRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -412,17 +414,21 @@ namespace LLaDAImage { GGML_ASSERT(!context_tensor.empty()); ggml_tensor* context = make_input(context_tensor); - pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("llada_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_llada_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = llada_image.forward(&runner_ctx, x, timesteps, context, pe); @@ -459,17 +465,21 @@ namespace LLaDAImage { ggml_tensor* source = make_input(source_tensor); GGML_ASSERT(x->ne[3] == 1); - pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(context->ne[1]), - semantic != nullptr ? static_cast(semantic->ne[1]) : 0, - ZImage::SEQ_MULTI_OF, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("llada_image.pe1"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shape(semantic); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_llada_image_edit_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(context->ne[1]), + semantic != nullptr ? static_cast(semantic->ne[1]) : 0, + ZImage::SEQ_MULTI_OF, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); int64_t W = x->ne[0]; diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index e10d5e829a..32e38012bd 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -11,6 +11,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -37,12 +38,12 @@ namespace LTXV { return mod; } - __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* shift, ggml_tensor* scale) { - shift = align_token_modulation(ctx, x, shift); - scale = align_token_modulation(ctx, x, scale); + shift = align_token_modulation(ctx->ggml_ctx, x, shift); + scale = align_token_modulation(ctx->ggml_ctx, x, scale); return Flux::modulate(ctx, x, shift, scale, true); } @@ -597,9 +598,9 @@ namespace LTXV { : embedding_dim(embedding_dim), embedding_coefficient(embedding_coefficient) { blocks["emb.timestep_embedder"] = std::make_shared(embedding_dim); blocks["linear"] = std::make_shared(embedding_dim, - embedding_coefficient * embedding_dim, - true, - true); + embedding_coefficient * embedding_dim, + true, + true); } std::pair forward(GGMLRunnerContext* ctx, @@ -820,7 +821,7 @@ namespace LTXV { auto gate_mlp = mods[5]; auto x_norm = rms_norm(ctx->ggml_ctx, x); - x_norm = LTXV::modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa); + x_norm = LTXV::modulate(ctx, x_norm, shift_msa, scale_msa); auto msa = attn1->forward(ctx, x_norm, nullptr, self_attention_mask, pe); x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, msa, gate_msa)); @@ -830,12 +831,12 @@ namespace LTXV { auto gate_q = mods[8]; auto q = rms_norm(ctx->ggml_ctx, x); - q = LTXV::modulate(ctx->ggml_ctx, q, shift_q, scale_q); + q = LTXV::modulate(ctx, q, shift_q, scale_q); auto context_mod = context; if (prompt_timestep != nullptr) { auto prompt_mods = get_prompt_scale_shift_values(ctx, prompt_timestep); - context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, prompt_mods[0], prompt_mods[1]); + context_mod = LTXV::modulate(ctx, context_mod, prompt_mods[0], prompt_mods[1]); } auto mca = attn2->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr); @@ -846,7 +847,7 @@ namespace LTXV { } auto y = rms_norm(ctx->ggml_ctx, x); - y = LTXV::modulate(ctx->ggml_ctx, y, shift_mlp, scale_mlp); + y = LTXV::modulate(ctx, y, shift_mlp, scale_mlp); auto mlp_out = ff->forward(ctx, y); x = ggml_add(ctx->ggml_ctx, x, apply_gate(ctx->ggml_ctx, mlp_out, gate_mlp)); return x; @@ -1198,11 +1199,11 @@ namespace LTXV { if (cross_attention_adaln) { auto q_mods = get_ada_values(ctx, table, timestep, dim, 9, 6, 3); auto q = rms_norm(ctx->ggml_ctx, x); - q = LTXV::modulate(ctx->ggml_ctx, q, q_mods[0], q_mods[1]); + q = LTXV::modulate(ctx, q, q_mods[0], q_mods[1]); auto context_mod = context; if (prompt_timestep != nullptr && prompt_table != nullptr) { auto p_mods = get_ada_values(ctx, prompt_table, prompt_timestep, dim, 2); - context_mod = LTXV::modulate(ctx->ggml_ctx, context_mod, p_mods[0], p_mods[1]); + context_mod = LTXV::modulate(ctx, context_mod, p_mods[0], p_mods[1]); } auto out = attn->forward(ctx, q, context_mod, attention_mask, nullptr, nullptr); return apply_gate(ctx->ggml_ctx, out, q_mods[2]); @@ -1249,7 +1250,7 @@ namespace LTXV { auto v_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6); auto v_norm = rms_norm(ctx->ggml_ctx, vx); - v_norm = LTXV::modulate(ctx->ggml_ctx, v_norm, v_mods[0], v_mods[1]); + v_norm = LTXV::modulate(ctx, v_norm, v_mods[0], v_mods[1]); auto v_sa = attn1->forward(ctx, v_norm, nullptr, self_attention_mask, v_pe); vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_sa, v_mods[2])); auto v_txt = apply_text_cross_attention(ctx, @@ -1257,7 +1258,7 @@ namespace LTXV { v_context, attn2.get(), v_table, - cross_attention_adaln ? params["prompt_scale_shift_table"] : nullptr, + cross_attention_adaln ? params["prompt_scale_shift_table"] : nullptr, v_timestep, v_prompt_timestep, v_dim, @@ -1267,7 +1268,7 @@ namespace LTXV { if (run_ax) { auto a_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6); auto a_norm = rms_norm(ctx->ggml_ctx, ax); - a_norm = LTXV::modulate(ctx->ggml_ctx, a_norm, a_mods[0], a_mods[1]); + a_norm = LTXV::modulate(ctx, a_norm, a_mods[0], a_mods[1]); auto a_sa = audio_attn1->forward(ctx, a_norm, nullptr, nullptr, a_pe); ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_sa, a_mods[2])); auto a_txt = apply_text_cross_attention(ctx, @@ -1275,7 +1276,7 @@ namespace LTXV { a_context, audio_attn2.get(), a_table, - cross_attention_adaln ? params["audio_prompt_scale_shift_table"] : nullptr, + cross_attention_adaln ? params["audio_prompt_scale_shift_table"] : nullptr, a_timestep, a_prompt_timestep, a_dim, @@ -1290,8 +1291,8 @@ namespace LTXV { auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4); auto a2v_audio = get_ada_values(ctx, a2v_audio_table, a_cross_scale_shift_timestep, a_dim, 4); auto a2v_video = get_ada_values(ctx, a2v_video_table, v_cross_scale_shift_timestep, v_dim, 4); - auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, a2v_video[1], a2v_video[0]); - auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, a2v_audio[1], a2v_audio[0]); + auto vx_scaled = LTXV::modulate(ctx, vx_norm3, a2v_video[1], a2v_video[0]); + auto ax_scaled = LTXV::modulate(ctx, ax_norm3, a2v_audio[1], a2v_audio[0]); auto a2v_out = audio_to_video_attn->forward(ctx, vx_scaled, ax_scaled, nullptr, v_cross_pe, a_cross_pe); auto a2v_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 4, 5); auto a2v_gate = get_ada_values(ctx, a2v_gate_table, v_cross_gate_timestep, v_dim, 1)[0]; @@ -1303,8 +1304,8 @@ namespace LTXV { auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4); auto v2a_audio = get_ada_values(ctx, v2a_audio_table, a_cross_scale_shift_timestep, a_dim, 4); auto v2a_video = get_ada_values(ctx, v2a_video_table, v_cross_scale_shift_timestep, v_dim, 4); - auto ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_norm3, v2a_audio[3], v2a_audio[2]); - auto vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_norm3, v2a_video[3], v2a_video[2]); + auto ax_scaled = LTXV::modulate(ctx, ax_norm3, v2a_audio[3], v2a_audio[2]); + auto vx_scaled = LTXV::modulate(ctx, vx_norm3, v2a_video[3], v2a_video[2]); auto v2a_out = video_to_audio_attn->forward(ctx, ax_scaled, vx_scaled, nullptr, a_cross_pe, v_cross_pe); auto v2a_gate_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 4, 5); auto v2a_gate = get_ada_values(ctx, v2a_gate_table, a_cross_gate_timestep, a_dim, 1)[0]; @@ -1312,14 +1313,14 @@ namespace LTXV { } auto a_ff_mods = get_ada_values(ctx, a_table, a_timestep, a_dim, cross_attention_adaln ? 9 : 6, 3, 3); auto ax_scaled = rms_norm(ctx->ggml_ctx, ax); - ax_scaled = LTXV::modulate(ctx->ggml_ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]); + ax_scaled = LTXV::modulate(ctx, ax_scaled, a_ff_mods[0], a_ff_mods[1]); auto a_ff_out = audio_ff->forward(ctx, ax_scaled); ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, a_ff_out, a_ff_mods[2])); } auto v_ff_mods = get_ada_values(ctx, v_table, v_timestep, v_dim, cross_attention_adaln ? 9 : 6, 3, 3); auto vx_scaled = rms_norm(ctx->ggml_ctx, vx); - vx_scaled = LTXV::modulate(ctx->ggml_ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]); + vx_scaled = LTXV::modulate(ctx, vx_scaled, v_ff_mods[0], v_ff_mods[1]); auto v_ff_out = ff->forward(ctx, vx_scaled); vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, v_ff_out, v_ff_mods[2])); @@ -1719,14 +1720,14 @@ namespace LTXV { auto v_shift_scale = get_output_scale_shift(ctx, params["scale_shift_table"], v_embedded_time, config.hidden_size); vx = norm_out->forward(ctx, vx); - vx = LTXV::modulate(ctx->ggml_ctx, vx, v_shift_scale[0], v_shift_scale[1]); + vx = LTXV::modulate(ctx, vx, v_shift_scale[0], v_shift_scale[1]); vx = proj_out->forward(ctx, vx); vx = unpatchify_video(ctx, vx, width, height, frames); if (ax != nullptr && audio_time > 0) { auto a_shift_scale = get_output_scale_shift(ctx, params["audio_scale_shift_table"], a_embedded_time, config.audio_hidden_size); ax = audio_norm_out->forward(ctx, ax); - ax = LTXV::modulate(ctx->ggml_ctx, ax, a_shift_scale[0], a_shift_scale[1]); + ax = LTXV::modulate(ctx, ax, a_shift_scale[0], a_shift_scale[1]); ax = audio_proj_out->forward(ctx, ax); ax = unpatchify_audio(ctx, ax, audio_time); } @@ -1736,14 +1737,9 @@ namespace LTXV { }; struct LTXAVRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; LTXAVConfig config; LTXAVModelBlock model; - std::vector video_pe_vec; - std::vector audio_pe_vec; - std::vector video_cross_pe_vec; - std::vector audio_cross_pe_vec; - std::vector connector_pe_vec; - std::vector audio_connector_pe_vec; sd::Tensor vx_input_cache; sd::Tensor ax_input_cache; @@ -1853,79 +1849,101 @@ namespace LTXV { float video_frame_rate = frame_rate > 0.f ? frame_rate : 24.f; int64_t video_token_count = vx->ne[0] * vx->ne[1] * vx->ne[2]; bool has_video_positions = !video_positions_tensor.empty(); - if (has_video_positions) { - GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count); - video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor, - static_cast(config.hidden_size), - static_cast(config.num_attention_heads), - config.positional_embedding_theta, - config.positional_embedding_max_pos, - config.use_middle_indices_grid); - } else { - video_pe_vec = build_video_rope_matrix(vx->ne[0], - vx->ne[1], - vx->ne[2], - static_cast(config.hidden_size), - static_cast(config.num_attention_heads), - video_frame_rate, - config.positional_embedding_theta, - config.positional_embedding_max_pos, - config.vae_scale_factors, - config.causal_temporal_positioning, - config.use_middle_indices_grid); - } - auto video_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads); - ggml_set_name(video_pe, "ltxav_video_pe"); - set_backend_tensor_data(video_pe, video_pe_vec.data()); + sd::PECacheKey video_pe_key("ltxav_video_pe"); + video_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor).add(config.positional_embedding_theta); + video_pe_key.add(config.causal_temporal_positioning).add(config.use_middle_indices_grid).add(config.hidden_size); + video_pe_key.add(config.num_attention_heads).add(config.attention_head_dim).add(config.positional_embedding_max_pos); + video_pe_key.add(std::get<0>(config.vae_scale_factors)).add(std::get<1>(config.vae_scale_factors)); + video_pe_key.add(std::get<2>(config.vae_scale_factors)); + auto video_pe = make_input(pe_cache_.get_or_create(video_pe_key, [&]() { + std::vector video_pe_vec; + if (has_video_positions) { + GGML_ASSERT(video_positions_tensor.shape()[2] == video_token_count); + video_pe_vec = build_video_rope_matrix_from_positions(video_positions_tensor, + static_cast(config.hidden_size), + static_cast(config.num_attention_heads), + config.positional_embedding_theta, + config.positional_embedding_max_pos, + config.use_middle_indices_grid); + } else { + video_pe_vec = build_video_rope_matrix(vx->ne[0], + vx->ne[1], + vx->ne[2], + static_cast(config.hidden_size), + static_cast(config.num_attention_heads), + video_frame_rate, + config.positional_embedding_theta, + config.positional_embedding_max_pos, + config.vae_scale_factors, + config.causal_temporal_positioning, + config.use_middle_indices_grid); + } + return sd::Tensor({2, 2, config.attention_head_dim / 2, video_token_count * config.num_attention_heads}, std::move(video_pe_vec)); + })); ggml_tensor* audio_pe = nullptr; ggml_tensor* video_cross_pe = nullptr; ggml_tensor* audio_cross_pe = nullptr; if (ax != nullptr && ggml_nelements(ax) > 0 && ax->ne[1] > 0) { - audio_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(config.audio_hidden_size), - static_cast(config.audio_num_attention_heads), - config.positional_embedding_theta, - config.audio_positional_embedding_max_pos[0], - config.use_middle_indices_grid); - audio_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads); - ggml_set_name(audio_pe, "ltxav_audio_pe"); - set_backend_tensor_data(audio_pe, audio_pe_vec.data()); + sd::PECacheKey audio_pe_key("ltxav_audio_pe"); + audio_pe_key.add(ax->ne[1]).add(config.audio_hidden_size).add(config.audio_num_attention_heads); + audio_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta); + audio_pe_key.add(config.audio_positional_embedding_max_pos[0]).add(config.use_middle_indices_grid); + audio_pe = make_input(pe_cache_.get_or_create(audio_pe_key, [&]() { + auto audio_pe_vec = build_audio_rope_matrix(ax->ne[1], + static_cast(config.audio_hidden_size), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + config.audio_positional_embedding_max_pos[0], + config.use_middle_indices_grid); + return sd::Tensor({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_pe_vec)); + })); int temporal_max_pos = std::max(config.positional_embedding_max_pos[0], config.audio_positional_embedding_max_pos[0]); - if (has_video_positions) { - video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor, - static_cast(config.audio_cross_attention_dim), - static_cast(config.audio_num_attention_heads), - config.positional_embedding_theta, - temporal_max_pos, - true); - } else { - video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0], - vx->ne[1], - vx->ne[2], - static_cast(config.audio_cross_attention_dim), - static_cast(config.audio_num_attention_heads), - video_frame_rate, - config.positional_embedding_theta, - temporal_max_pos, - std::get<0>(config.vae_scale_factors), - config.causal_temporal_positioning, - true); - } - video_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads); - ggml_set_name(video_cross_pe, "ltxav_video_cross_pe"); - set_backend_tensor_data(video_cross_pe, video_cross_pe_vec.data()); - - audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1], - static_cast(config.audio_cross_attention_dim), - static_cast(config.audio_num_attention_heads), - config.positional_embedding_theta, - temporal_max_pos, - true); - audio_cross_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads); - ggml_set_name(audio_cross_pe, "ltxav_audio_cross_pe"); - set_backend_tensor_data(audio_cross_pe, audio_cross_pe_vec.data()); + sd::PECacheKey video_cross_pe_key("ltxav_video_cross_pe"); + video_cross_pe_key.add_shape(vx).add(video_frame_rate).add(video_positions_tensor); + video_cross_pe_key.add(config.positional_embedding_theta).add(config.causal_temporal_positioning); + video_cross_pe_key.add(config.use_middle_indices_grid).add(config.audio_cross_attention_dim); + video_cross_pe_key.add(config.audio_num_attention_heads).add(config.audio_attention_head_dim); + video_cross_pe_key.add(temporal_max_pos).add(std::get<0>(config.vae_scale_factors)); + video_cross_pe = make_input(pe_cache_.get_or_create(video_cross_pe_key, [&]() { + std::vector video_cross_pe_vec; + if (has_video_positions) { + video_cross_pe_vec = build_video_temporal_rope_matrix_from_positions(video_positions_tensor, + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + temporal_max_pos, + true); + } else { + video_cross_pe_vec = build_video_temporal_rope_matrix(vx->ne[0], + vx->ne[1], + vx->ne[2], + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + video_frame_rate, + config.positional_embedding_theta, + temporal_max_pos, + std::get<0>(config.vae_scale_factors), + config.causal_temporal_positioning, + true); + } + return sd::Tensor({2, 2, config.audio_attention_head_dim / 2, video_token_count * config.audio_num_attention_heads}, std::move(video_cross_pe_vec)); + })); + + sd::PECacheKey audio_cross_pe_key("ltxav_audio_cross_pe"); + audio_cross_pe_key.add(ax->ne[1]).add(config.audio_cross_attention_dim).add(config.audio_num_attention_heads); + audio_cross_pe_key.add(config.audio_attention_head_dim).add(config.positional_embedding_theta); + audio_cross_pe_key.add(temporal_max_pos); + audio_cross_pe = make_input(pe_cache_.get_or_create(audio_cross_pe_key, [&]() { + auto audio_cross_pe_vec = build_audio_rope_matrix(ax->ne[1], + static_cast(config.audio_cross_attention_dim), + static_cast(config.audio_num_attention_heads), + config.positional_embedding_theta, + temporal_max_pos, + true); + return sd::Tensor({2, 2, config.audio_attention_head_dim / 2, ax->ne[1] * config.audio_num_attention_heads}, std::move(audio_cross_pe_vec)); + })); } bool needs_video_connector_pe = @@ -1941,10 +1959,13 @@ namespace LTXV { int64_t target_len = std::max(1024, seq_len); int64_t duplications = (target_len + config.connector_num_registers - 1) / config.connector_num_registers; int64_t full_len = seq_len + duplications * config.connector_num_registers - seq_len; - connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.connector_hidden_size), static_cast(config.connector_num_heads), 10000.f, 4096.f, true); - video_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads); - ggml_set_name(video_connector_pe, "ltxav_video_connector_pe"); - set_backend_tensor_data(video_connector_pe, connector_pe_vec.data()); + sd::PECacheKey video_connector_pe_key("ltxav_video_connector_pe"); + video_connector_pe_key.add(full_len).add(config.connector_hidden_size).add(config.connector_num_heads); + video_connector_pe_key.add(config.connector_head_dim); + video_connector_pe = make_input(pe_cache_.get_or_create(video_connector_pe_key, [&]() { + auto connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.connector_hidden_size), static_cast(config.connector_num_heads), 10000.f, 4096.f, true); + return sd::Tensor({2, 2, config.connector_head_dim / 2, full_len * config.connector_num_heads}, std::move(connector_pe_vec)); + })); } bool run_audio_context = @@ -1961,14 +1982,17 @@ namespace LTXV { context->ne[1] < 1024)); ggml_tensor* audio_connector_pe = nullptr; if (needs_audio_connector_pe) { - int64_t seq_len = context->ne[1]; - int64_t target_len = std::max(1024, seq_len); - int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers; - int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len; - audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.audio_connector_hidden_size), static_cast(config.audio_connector_num_heads), 10000.f, 4096.f, true); - audio_connector_pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads); - ggml_set_name(audio_connector_pe, "ltxav_audio_connector_pe"); - set_backend_tensor_data(audio_connector_pe, audio_connector_pe_vec.data()); + int64_t seq_len = context->ne[1]; + int64_t target_len = std::max(1024, seq_len); + int64_t duplications = (target_len + config.audio_connector_num_registers - 1) / config.audio_connector_num_registers; + int64_t full_len = seq_len + duplications * config.audio_connector_num_registers - seq_len; + sd::PECacheKey audio_connector_pe_key("ltxav_audio_connector_pe"); + audio_connector_pe_key.add(full_len).add(config.audio_connector_hidden_size).add(config.audio_connector_num_heads); + audio_connector_pe_key.add(config.audio_connector_head_dim); + audio_connector_pe = make_input(pe_cache_.get_or_create(audio_connector_pe_key, [&]() { + auto audio_connector_pe_vec = build_1d_rope_matrix(full_len, static_cast(config.audio_connector_hidden_size), static_cast(config.audio_connector_num_heads), 10000.f, 4096.f, true); + return sd::Tensor({2, 2, config.audio_connector_head_dim / 2, full_len * config.audio_connector_num_heads}, std::move(audio_connector_pe_vec)); + })); } auto runner_ctx = get_context(); diff --git a/src/model/diffusion/mage_flow.hpp b/src/model/diffusion/mage_flow.hpp index 1b9b9f9be9..415c0f7228 100644 --- a/src/model/diffusion/mage_flow.hpp +++ b/src/model/diffusion/mage_flow.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "model/diffusion/qwen_image.hpp" @@ -57,11 +58,12 @@ namespace MageFlow { }; struct MageFlowRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: Qwen::QwenImageConfig config; Qwen::QwenImageModel mage_flow; std::shared_ptr time_text_embed; - std::vector pe_vec; MageFlowRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -110,16 +112,20 @@ namespace MageFlow { } int batch_size = static_cast(x->ne[3]); - pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - batch_size, - static_cast(context->ne[1]), - ref_latents, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("mage_flow.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(batch_size).add_shapes(ref_latents); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_mage_flow_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + batch_size, + static_cast(context->ne[1]), + ref_latents, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); time_text_embed->prepare(timesteps_tensor); auto runner_ctx = get_context(); diff --git a/src/model/diffusion/ming_image.hpp b/src/model/diffusion/ming_image.hpp index f19be99033..1de9ac157e 100644 --- a/src/model/diffusion/ming_image.hpp +++ b/src/model/diffusion/ming_image.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ #define __SD_MODEL_DIFFUSION_MING_IMAGE_HPP__ +#include "core/pe_cache.h" #include "z_image.hpp" namespace MingImage { @@ -72,9 +73,9 @@ namespace MingImage { }; struct MingImageRunner : DiffusionModelRunner { + sd::PECache pe_cache_; MingImageConfig config; MingImageModel model; - std::vector pe_values; MingImageRunner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix, std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), @@ -109,19 +110,23 @@ namespace MingImage { const int64_t n_txt = context->ne[1] + direct->ne[1]; const int64_t n_img = ((x->ne[0] + config.patch_size - 1) / config.patch_size) * ((x->ne[1] + config.patch_size - 1) / config.patch_size); - auto padded = finish_rope_pe(Rope::gen_z_image_pe( - static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, - static_cast(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED, - config.theta, config.axes_dim)); - // Zero-masked alignment tokens cannot affect valid queries. Omit them while - // retaining the padded caption length used to position image tokens. - const size_t stride = config.axes_dim_sum * 2; - const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); - pe_values.assign(padded.begin(), padded.begin() + n_txt * stride); - pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride, - padded.begin() + (padded_txt + n_img) * stride); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, n_txt + n_img); - set_backend_tensor_data(pe, pe_values.data()); + sd::PECacheKey pe_key("ming_image.pe"); + pe_key.add_shape(x).add(n_txt).add(config.patch_size).add(config.theta).add(config.axes_dim); + pe_key.add(config.axes_dim_sum).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto padded = finish_rope_pe(Rope::gen_z_image_pe( + static_cast(x->ne[1]), static_cast(x->ne[0]), config.patch_size, 1, + static_cast(n_txt), ZImage::SEQ_MULTI_OF, {}, Rope::RefIndexMode::FIXED, + config.theta, config.axes_dim)); + // Zero-masked alignment tokens cannot affect valid queries. Omit them while + // retaining the padded caption length used to position image tokens. + const size_t stride = config.axes_dim_sum * 2; + const int64_t padded_txt = n_txt + Rope::bound_mod(static_cast(n_txt), ZImage::SEQ_MULTI_OF); + std::vector pe_values(padded.begin(), padded.begin() + n_txt * stride); + pe_values.insert(pe_values.end(), padded.begin() + padded_txt * stride, + padded.begin() + (padded_txt + n_img) * stride); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, n_txt + n_img}, std::move(pe_values)); + })); auto ctx = get_context(); auto out = model.forward(&ctx, x, t, context, direct, pe); ggml_build_forward_expand(gf, out); diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 7caaeef7c2..d25f2a6dbe 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -13,6 +13,7 @@ #include "core/ggml_tensor_utils.h" #include "core/ggml_graph_cut.h" +#include "model/common/block.hpp" #include "model/diffusion/dit.hpp" #include "model/diffusion/model.hpp" @@ -188,8 +189,11 @@ namespace MiniMaxH3 { }; struct MLP : public UnaryBlock { + int64_t ffn_hidden_size; + MLP(int64_t hidden_size, - int64_t ffn_hidden_size) { + int64_t ffn_hidden_size) + : ffn_hidden_size(ffn_hidden_size) { blocks["fc1"] = std::make_shared(hidden_size, ffn_hidden_size * 2, false, false, true, 1.f / 128.f); blocks["fc2"] = std::make_shared(ffn_hidden_size, hidden_size, false, false, true, 1.f / 128.f); } @@ -231,23 +235,51 @@ namespace MiniMaxH3 { #endif } + void init_params(ggml_context* ctx, + const String2TensorStorage& tensors = {}, + const std::string prefix = "") override { + // INT8 accumulates in INT32 and returns FP32, so it needs no FP16 overflow scaling. + for (const auto* name : {"fc1", "fc2"}) { + auto linear = std::dynamic_pointer_cast(blocks[name]); + linear->set_scale(get_type(prefix + name + ".weight", tensors, GGML_TYPE_F32) == GGML_TYPE_I8 + ? 1.f + : 1.f / 128.f); + } + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); - auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); - auto h = fc1->forward(ctx, x); - if (fast_mlp() && ctx->backend != nullptr && h->type == GGML_TYPE_F32 && ggml_is_contiguous(h)) { - // silu(u) * v straight from the two halves of fc1's output: same per-element - // silu and product as the copied-chunk graph, without the two chunk copies. - auto halves = ggml_ext_chunk(ctx->ggml_ctx, h, 2, 0, false); - auto gated = ggml_swiglu_split(ctx->ggml_ctx, halves[0], halves[1]); - if (ggml_backend_supports_op(ctx->backend, gated)) { - return fc2->forward(ctx, gated); + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + auto forward_chunk = [&](ggml_tensor* input) { + auto projected = fc1->forward(ctx, input); + auto activated = ggml_swiglu(ctx->ggml_ctx, projected); + return fc2->forward(ctx, activated); + }; + + int64_t chunk_size = x->ne[1]; + if (ctx->backend != nullptr) { + auto buft = ggml_backend_get_default_buffer_type(ctx->backend); + size_t alignment = ggml_backend_buft_get_alignment(buft); + size_t max_size = ggml_backend_buft_get_max_size(buft); + max_size -= max_size % alignment; + // Matmul outputs are F32 even for quantized weights. Split before fc1 + // so its expanded activation never needs an oversized backend buffer. + size_t token_bytes = ggml_row_size(GGML_TYPE_F32, std::max(x->ne[0], 2 * ffn_hidden_size)) * x->ne[2] * x->ne[3]; + chunk_size = static_cast(std::min(chunk_size, std::max(1, max_size / token_bytes))); + } + if (chunk_size == x->ne[1]) { + return forward_chunk(x); + } + + std::vector chunks; + for (int64_t start = 0; start < x->ne[1]; start += chunk_size) { + auto input = ggml_ext_slice(ctx->ggml_ctx, x, 1, start, std::min(start + chunk_size, x->ne[1]), false); + if (!ggml_is_contiguous(input)) { + input = ggml_cont(ctx->ggml_ctx, input); } + chunks.push_back(forward_chunk(input)); } - auto uv = ggml_ext_chunk(ctx->ggml_ctx, h, 2, 0); - return fc2->forward(ctx, ggml_mul(ctx->ggml_ctx, - ggml_silu(ctx->ggml_ctx, uv[0]), - uv[1])); + return ggml_ext_vec_concat(ctx->ggml_ctx, chunks, 1); } }; @@ -262,6 +294,14 @@ namespace MiniMaxH3 { auto ctx = runner_ctx->ggml_ctx; int64_t rot_dim = pe->ne[2] * 2; GGML_ASSERT(rot_dim <= x->ne[0]); +#ifndef SD_USE_UPSTREAM_GGML + if (runner_ctx->backend != nullptr && x->type == GGML_TYPE_F32 && pe->type == GGML_TYPE_F32) { + auto out = sd_ggml_rope_apply(ctx, runner_ctx->backend, x, pe, false); + if (out && ggml_backend_supports_op(runner_ctx->backend, out)) { + return out; + } + } +#endif auto rotated = Rope::apply_rope(runner_ctx, ggml_ext_slice(ctx, x, 0, 0, rot_dim), pe, @@ -363,29 +403,31 @@ namespace MiniMaxH3 { #endif } - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* pe = nullptr) { - auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv_proj"]); - auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); - auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); - auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); - - int64_t sequence = x->ne[1]; - int64_t batch = x->ne[2] * x->ne[3]; - auto qkv_out = qkv_proj->forward(ctx, x); - if (fast_qkv()) { - auto out = forward_head_major(ctx, qkv_out, sequence, batch, pe); - if (out != nullptr) { - return out_proj->forward(ctx, out); - } - } - auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_out, 3, 0); - auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], head_dim, heads, sequence, batch); - auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv[1], head_dim, heads, sequence, batch); - auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv[2], head_dim, heads, sequence, batch); - q = q_norm->forward(ctx, q); - k = k_norm->forward(ctx, k); + ggml_tensor* project_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) { + return std::dynamic_pointer_cast(blocks["qkv_proj"])->forward(ctx, x); + } + + ggml_tensor* project_output(GGMLRunnerContext* ctx, ggml_tensor* x) { + return std::dynamic_pointer_cast(blocks["out_proj"])->forward(ctx, x); + } + + ggml_tensor* attend(GGMLRunnerContext* ctx, + ggml_tensor* projected, + ggml_tensor* pe = nullptr) { + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + int64_t sequence = projected->ne[1]; + int64_t batch = projected->ne[2] * projected->ne[3]; + auto view = [&](int index) { + return ggml_view_4d(ctx->ggml_ctx, projected, head_dim, heads, sequence, batch, + head_dim * projected->nb[0], projected->nb[1], projected->nb[2], + index * heads * head_dim * projected->nb[0]); + }; + auto q = view(0); + auto k = view(1); + auto v = view(2); + q = q_norm->forward(ctx, q); + k = k_norm->forward(ctx, k); if (pe != nullptr) { q = apply_partial_rope(ctx, q, pe); k = apply_partial_rope(ctx, k, pe); @@ -402,7 +444,20 @@ namespace MiniMaxH3 { true, ctx->flash_attn_enabled, 1.f / 128.f); - return out_proj->forward(ctx, out); + return out; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe = nullptr) { + auto qkv_out = project_qkv(ctx, x); + if (fast_qkv()) { + auto out = forward_head_major(ctx, qkv_out, x->ne[1], x->ne[2] * x->ne[3], pe); + if (out != nullptr) { + return project_output(ctx, out); + } + } + return project_output(ctx, attend(ctx, qkv_out, pe)); } }; @@ -411,9 +466,9 @@ namespace MiniMaxH3 { blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["attn"] = std::make_shared(config.hidden_size, - config.num_attention_heads, - config.attention_head_dim, - config.qk_norm_eps); + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps); blocks["mlp"] = std::make_shared(config.hidden_size, config.ffn_hidden_size); } @@ -516,50 +571,6 @@ namespace MiniMaxH3 { return ggml_ext_chunk(ctx, selected, expand, 0, !fast_views()); } - static ggml_tensor* modulate_segments(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* projection, - const std::vector& segments, - int64_t hidden_size, - int expand, - int modalities, - int shift_index, - int scale_index) { - ggml_tensor* out = nullptr; - for (const auto& segment : segments) { - auto mods = modulation_row(ctx, - projection, - hidden_size, - expand, - modalities, - segment.modulation_row); - auto part = ggml_ext_slice(ctx, x, 1, segment.start, segment.end, !fast_views()); - part = ggml_add(ctx, - ggml_add(ctx, part, ggml_mul(ctx, part, mods[scale_index])), - mods[shift_index]); - out = out == nullptr ? part : ggml_concat(ctx, out, part, 1); - } - return out; - } - - static ggml_tensor* gated_residual_segments(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* update, - ggml_tensor* projection, - const std::vector& segments, - int64_t hidden_size, - int gate_index) { - ggml_tensor* out = nullptr; - for (const auto& segment : segments) { - auto mods = modulation_row(ctx, projection, hidden_size, 6, 3, segment.modulation_row); - auto base = ggml_ext_slice(ctx, x, 1, segment.start, segment.end, !fast_views()); - auto add = ggml_ext_slice(ctx, update, 1, segment.start, segment.end, !fast_views()); - auto part = ggml_add(ctx, base, ggml_mul(ctx, add, mods[gate_index])); - out = out == nullptr ? part : ggml_concat(ctx, out, part, 1); - } - return out; - } - struct TransformerBlock : public GGMLBlock { Config config; @@ -568,11 +579,11 @@ namespace MiniMaxH3 { blocks["norm1"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["norm2"] = std::make_shared(config.hidden_size, config.norm_eps); blocks["attn"] = std::make_shared(config.hidden_size, - config.num_attention_heads, - config.attention_head_dim, - config.qk_norm_eps); + config.num_attention_heads, + config.attention_head_dim, + config.qk_norm_eps); blocks["mlp"] = std::make_shared(config.hidden_size, - config.ffn_hidden_size); + config.ffn_hidden_size); blocks["adaln_proj"] = std::make_shared(config.time_embed_dim, config.hidden_size, 6, @@ -677,38 +688,41 @@ namespace MiniMaxH3 { auto adaln = std::dynamic_pointer_cast(blocks["adaln_proj"]); auto mods = adaln->forward(ctx, t_emb); - auto h = modulate_segments(ctx->ggml_ctx, - norm1->forward(ctx, x), - mods, - segments, - config.hidden_size, - 6, - 3, - 0, - 1); - x = gated_residual_segments(ctx->ggml_ctx, - x, - attn->forward(ctx, h, pe), - mods, - segments, - config.hidden_size, - 2); - h = modulate_segments(ctx->ggml_ctx, - norm2->forward(ctx, x), - mods, - segments, - config.hidden_size, - 6, - 3, - 3, - 4); - return gated_residual_segments(ctx->ggml_ctx, - x, - mlp->forward(ctx, h), - mods, - segments, - config.hidden_size, - 5); + std::vector> rows; + std::vector inputs; + for (const auto& segment : segments) { + rows.push_back(modulation_row(ctx->ggml_ctx, mods, config.hidden_size, + 6, 3, segment.modulation_row)); + inputs.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end, false)); + } + auto modulate = [&](ggml_tensor* value, const std::vector& row, int shift, int scale) { + return ggml_ext_modulate(ctx, value, row[shift], row[scale]); + }; + ggml_tensor* qkv = nullptr; + for (size_t i = 0; i < segments.size(); ++i) { + auto h = modulate(norm1->forward(ctx, inputs[i]), rows[i], 0, 1); + h = attn->project_qkv(ctx, h); + if (qkv == nullptr) { + qkv = h; + continue; + } + qkv = ggml_concat(ctx->ggml_ctx, qkv, h, 1); + } + auto attended = attn->attend(ctx, qkv, pe); + std::vector projections; + for (const auto& segment : segments) { + auto part = ggml_ext_slice(ctx->ggml_ctx, attended, 1, segment.start, segment.end, false); + projections.push_back(attn->project_output(ctx, part)); + } + ggml_tensor* output = nullptr; + for (size_t i = 0; i < segments.size(); ++i) { + auto residual = ggml_add(ctx->ggml_ctx, inputs[i], + ggml_mul(ctx->ggml_ctx, projections[i], rows[i][2])); + auto h = modulate(norm2->forward(ctx, residual), rows[i], 3, 4); + h = ggml_add(ctx->ggml_ctx, residual, ggml_mul(ctx->ggml_ctx, mlp->forward(ctx, h), rows[i][5])); + output = output == nullptr ? h : ggml_concat(ctx->ggml_ctx, output, h, 1); + } + return output; } }; @@ -743,8 +757,8 @@ namespace MiniMaxH3 { auto row = modulation_row(ctx->ggml_ctx, mods, config.hidden_size, 2, 1, segment.modulation_row); auto value = norm->forward(ctx, ggml_ext_slice(ctx->ggml_ctx, x, 1, segment.start, segment.end)); return ggml_add(ctx->ggml_ctx, - ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])), - row[0]); + ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[1])), + row[0]); }; return {video_out->forward(ctx, apply(video)), audio_out->forward(ctx, apply(audio))}; @@ -1107,10 +1121,10 @@ namespace MiniMaxH3 { const auto& condition = condition_videos[index]; auto axes = condition_spatial_axes(condition); int64_t count = condition.shape()[2] * - static_cast(axes.first.size() * axes.second.size()); - bool is_first = keyframe_indices.empty() || keyframe_indices[static_cast(index)] == 0; - float keyframe_t = is_first ? static_cast(text_len) - : static_cast(text_len) + video_duration - FRAME_RESCALE; + static_cast(axes.first.size() * axes.second.size()); + bool is_first = keyframe_indices.empty() || keyframe_indices[static_cast(index)] == 0; + float keyframe_t = is_first ? static_cast(text_len) + : static_cast(text_len) + video_duration - FRAME_RESCALE; for (int64_t t = 0; t < condition.shape()[2]; ++t) { for (float h : axes.first) { for (float w : axes.second) { @@ -1164,9 +1178,9 @@ namespace MiniMaxH3 { if (block.kind != MiniMaxH3ReferenceKind::AUDIO) { GGML_ASSERT(ref_video != nullptr); - auto axes = condition_spatial_axes(*ref_video); - int64_t count = ref_video->shape()[2] * - static_cast(axes.first.size() * axes.second.size()); + auto axes = condition_spatial_axes(*ref_video); + int64_t count = ref_video->shape()[2] * + static_cast(axes.first.size() * axes.second.size()); float video_end = append_video_positions(*ref_video, cursor); layout.sequence_segments.push_back({row, row + count, @@ -1216,8 +1230,6 @@ namespace MiniMaxH3 { struct MiniMaxH3Runner : public DiffusionModelRunner { Config config; MiniMaxH3Transformer3DModel model; - sd::Tensor video_input_cache; - sd::Tensor audio_input_cache; sd::Tensor position_input_cache; sd::Tensor timestep_feature_input_cache; sd::Tensor curve_index_input_cache; @@ -1243,29 +1255,6 @@ namespace MiniMaxH3 { model.get_param_tensors(tensors, prefix); } - std::pair, sd::Tensor> split_av_latents(const sd::Tensor& packed, - int audio_length) const { - GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); - int64_t spatial = packed.shape()[0] * packed.shape()[1] * packed.shape()[2]; - int64_t video_values = spatial * config.video_latent_channels; - sd::Tensor video({packed.shape()[0], - packed.shape()[1], - packed.shape()[2], - config.video_latent_channels, - 1}); - std::copy_n(packed.data(), static_cast(video_values), video.data()); - if (audio_length <= 0) { - return {video, {}}; - } - int64_t audio_values = audio_length * 2 * config.audio_latent_channels; - GGML_ASSERT(packed.numel() >= video_values + audio_values); - sd::Tensor audio({audio_length, 2, config.audio_latent_channels, 1}); - std::copy_n(packed.data() + video_values, - static_cast(audio_values), - audio.data()); - return {video, audio}; - } - ggml_tensor* merge_av_latents(ggml_context* ctx, ggml_tensor* video, ggml_tensor* audio) const { @@ -1296,14 +1285,26 @@ namespace MiniMaxH3 { int audio_length, float video_shift, float audio_shift) { - auto split = split_av_latents(packed, audio_length); - video_input_cache = std::move(split.first); - audio_input_cache = std::move(split.second); - GGML_ASSERT(!audio_input_cache.empty()); + GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); + GGML_ASSERT(audio_length > 0); GGML_ASSERT(!context_tensor.empty()); - - auto video = make_input(video_input_cache); - auto audio_carrier = make_input(audio_input_cache); + const int64_t spatial = packed.shape()[0] * packed.shape()[1] * packed.shape()[2]; + const int64_t video_values = spatial * config.video_latent_channels; + const int64_t audio_values = int64_t(audio_length) * 2 * config.audio_latent_channels; + GGML_ASSERT(packed.numel() >= video_values + audio_values); + auto packed_input = make_input(packed); + auto video = ggml_view_4d(compute_ctx, packed_input, + packed.shape()[0], packed.shape()[1], packed.shape()[2], + config.video_latent_channels, + packed.shape()[0] * sizeof(float), + packed.shape()[0] * packed.shape()[1] * sizeof(float), + spatial * sizeof(float), 0); + auto audio_carrier = ggml_view_4d(compute_ctx, packed_input, + audio_length, 2, config.audio_latent_channels, 1, + audio_length * sizeof(float), + int64_t(audio_length) * 2 * sizeof(float), + audio_values * sizeof(float), + video_values * sizeof(float)); auto context = make_input(context_tensor); std::vector condition_inputs; condition_inputs.reserve(condition_videos.size()); @@ -1325,9 +1326,9 @@ namespace MiniMaxH3 { // latent follows one sigma schedule. Restore x_a for the H3 network. auto audio = ggml_ext_scale(compute_ctx, audio_carrier, sigma_a / sigma_v); auto layout = build_layout(context_tensor.shape()[1], - video_input_cache.shape()[2], - video_input_cache.shape()[1], - video_input_cache.shape()[0], + packed.shape()[2], + packed.shape()[1], + packed.shape()[0], audio_length, condition_videos, condition_audios, diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index 0ddacf02eb..c6e90dae1c 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -9,6 +9,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -206,9 +207,9 @@ namespace MiniT2I { auto w1 = std::dynamic_pointer_cast(blocks["w1"]); auto w3 = std::dynamic_pointer_cast(blocks["w3"]); auto w2 = std::dynamic_pointer_cast(blocks["w2"]); - auto gate = ggml_silu(ctx->ggml_ctx, w1->forward(ctx, x)); + auto gate = w1->forward(ctx, x); auto up = w3->forward(ctx, x); - return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, up)); + return w2->forward(ctx, ggml_swiglu_split(ctx->ggml_ctx, gate, up)); } }; @@ -473,20 +474,9 @@ namespace MiniT2I { }; struct MiniT2IRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; MiniT2IConfig config; MMJiT model; - ggml_context* position_cache_ctx = nullptr; - ggml_backend_buffer_t position_cache_buffer = nullptr; - ggml_tensor* cached_pos_embed = nullptr; - ggml_tensor* cached_txt_pe = nullptr; - ggml_tensor* cached_joint_pe = nullptr; - int64_t cached_img_side = -1; - int64_t cached_txt_len = -1; - int64_t cached_hidden_size = -1; - int64_t cached_head_dim = -1; - bool cached_circular_x = false; - bool cached_circular_y = false; - MiniT2IRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "", @@ -497,10 +487,6 @@ namespace MiniT2I { model.init(params_ctx, tensor_storage_map, this->prefix); } - ~MiniT2IRunner() override { - free_position_cache(); - } - std::string get_desc() override { return "MiniT2I"; } @@ -509,75 +495,6 @@ namespace MiniT2I { model.get_param_tensors(tensors, prefix); } - void free_position_cache() { - if (position_cache_buffer != nullptr) { - ggml_backend_buffer_free(position_cache_buffer); - position_cache_buffer = nullptr; - } - if (position_cache_ctx != nullptr) { - ggml_free(position_cache_ctx); - position_cache_ctx = nullptr; - } - cached_pos_embed = nullptr; - cached_txt_pe = nullptr; - cached_joint_pe = nullptr; - cached_img_side = -1; - cached_txt_len = -1; - cached_hidden_size = -1; - cached_head_dim = -1; - } - - void ensure_position_cache(int64_t img_side, int64_t txt_len) { - if (cached_img_side == img_side && - cached_txt_len == txt_len && - cached_hidden_size == config.hidden_size && - cached_head_dim == config.head_dim && - cached_circular_x == circular_x_enabled && - cached_circular_y == circular_y_enabled && - cached_pos_embed != nullptr && - cached_txt_pe != nullptr && - cached_joint_pe != nullptr) { - return; - } - - free_position_cache(); - - auto pos_embed_vec = make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size)); - auto txt_pe_vec = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); - auto img_pe_vec = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); - auto joint_pe_vec = txt_pe_vec; - joint_pe_vec.insert(joint_pe_vec.end(), img_pe_vec.begin(), img_pe_vec.end()); - - ggml_init_params params; - params.mem_size = static_cast(3 * ggml_tensor_overhead()); - params.mem_buffer = nullptr; - params.no_alloc = true; - position_cache_ctx = ggml_init(params); - GGML_ASSERT(position_cache_ctx != nullptr); - - cached_pos_embed = ggml_new_tensor_3d(position_cache_ctx, GGML_TYPE_F32, config.hidden_size, img_side * img_side, 1); - ggml_set_name(cached_pos_embed, "minit2i.pos_embed"); - cached_txt_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len); - ggml_set_name(cached_txt_pe, "minit2i.txt_pe"); - cached_joint_pe = ggml_new_tensor_4d(position_cache_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, txt_len + img_side * img_side); - ggml_set_name(cached_joint_pe, "minit2i.joint_pe"); - - position_cache_buffer = ggml_backend_alloc_ctx_tensors(position_cache_ctx, runtime_backend); - GGML_ASSERT(position_cache_buffer != nullptr); - ggml_backend_buffer_set_usage(position_cache_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - ggml_backend_tensor_set(cached_pos_embed, pos_embed_vec.data(), 0, ggml_nbytes(cached_pos_embed)); - ggml_backend_tensor_set(cached_txt_pe, txt_pe_vec.data(), 0, ggml_nbytes(cached_txt_pe)); - ggml_backend_tensor_set(cached_joint_pe, joint_pe_vec.data(), 0, ggml_nbytes(cached_joint_pe)); - ggml_backend_synchronize(runtime_backend); - - cached_img_side = img_side; - cached_txt_len = txt_len; - cached_hidden_size = config.hidden_size; - cached_head_dim = config.head_dim; - cached_circular_x = circular_x_enabled; - cached_circular_y = circular_y_enabled; - } - ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, @@ -592,10 +509,29 @@ namespace MiniT2I { int64_t H = x->ne[1]; int64_t img_side = H / config.patch_size; int64_t txt_len = context->ne[1]; - ensure_position_cache(img_side, txt_len); + sd::PECacheKey pos_key("minit2i.pos_embed"); + pos_key.add(img_side).add(config.hidden_size); + auto pos_embed = make_input(pe_cache_.get_or_create(pos_key, [&]() { + return sd::Tensor({config.hidden_size, img_side * img_side, 1}, + make_2d_sincos_pos_embed(static_cast(img_side), static_cast(config.hidden_size))); + })); + sd::PECacheKey txt_key("minit2i.txt_pe"); + txt_key.add(txt_len).add(config.head_dim); + auto txt_pe = make_input(pe_cache_.get_or_create(txt_key, [&]() { + return sd::Tensor({2, 2, config.head_dim / 2, txt_len}, + make_text_rope(static_cast(txt_len), static_cast(config.head_dim))); + })); + sd::PECacheKey joint_key("minit2i.joint_pe"); + joint_key.add(img_side).add(txt_len).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled); + auto joint_pe = make_input(pe_cache_.get_or_create(joint_key, [&]() { + auto values = make_text_rope(static_cast(txt_len), static_cast(config.head_dim)); + auto image = finish_rope_pe(make_vision_rope(static_cast(img_side), static_cast(config.head_dim))); + values.insert(values.end(), image.begin(), image.end()); + return sd::Tensor({2, 2, config.head_dim / 2, txt_len + img_side * img_side}, std::move(values)); + })); auto runner_ctx = get_context(); - auto out = model.forward(&runner_ctx, x, context, mask, cached_pos_embed, cached_txt_pe, cached_joint_pe); + auto out = model.forward(&runner_ctx, x, context, mask, pos_embed, txt_pe, joint_pe); ggml_build_forward_expand(gf, out); return gf; } diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 42599c67a0..613cfc11a7 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -14,6 +14,7 @@ #include "core/util.h" #include "model/common/block.hpp" #include "model/common/ggml_block.hpp" +#include "model/diffusion/dit.hpp" #include "model/diffusion/model.hpp" #include "model_loader.h" @@ -371,18 +372,16 @@ class SelfAttention : public GGMLBlock { } }; -__STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, +__STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* shift, ggml_tensor* scale) { // x: [N, L, C] // scale: [N, C] // shift: [N, C] - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] - shift = ggml_reshape_3d(ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C] - x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); - x = ggml_add(ctx, x, shift); - return x; + scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] + shift = ggml_reshape_3d(ctx->ggml_ctx, shift, shift->ne[0], 1, shift->ne[1]); // [N, 1, C] + return ggml_ext_modulate(ctx, x, shift, scale); } struct DismantledBlock : public GGMLBlock { @@ -454,10 +453,10 @@ struct DismantledBlock : public GGMLBlock { auto x_norm = norm1->forward(ctx, x); - auto attn_in = modulate(ctx->ggml_ctx, x_norm, shift_msa, scale_msa); + auto attn_in = modulate(ctx, x_norm, shift_msa, scale_msa); auto qkv = attn->pre_attention(ctx, attn_in); - auto attn2_in = modulate(ctx->ggml_ctx, x_norm, shift_msa2, scale_msa2); + auto attn2_in = modulate(ctx, x_norm, shift_msa2, scale_msa2); auto qkv2 = attn2->pre_attention(ctx, attn2_in); return {qkv, qkv2, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp, gate_msa2}}; @@ -487,13 +486,13 @@ struct DismantledBlock : public GGMLBlock { auto scale_mlp = m_vec[4]; // [N, hidden_size] auto gate_mlp = m_vec[5]; // [N, hidden_size] - auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa); + auto attn_in = modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa); auto qkv = attn->pre_attention(ctx, attn_in); return {qkv, {x, gate_msa, shift_mlp, scale_mlp, gate_mlp}}; } else { - auto attn_in = modulate(ctx->ggml_ctx, norm1->forward(ctx, x), shift_msa, scale_msa); + auto attn_in = modulate(ctx, norm1->forward(ctx, x), shift_msa, scale_msa); auto qkv = attn->pre_attention(ctx, attn_in); return {qkv, {nullptr, nullptr, nullptr, nullptr, nullptr}}; @@ -532,7 +531,7 @@ struct DismantledBlock : public GGMLBlock { x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa)); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn2_out, gate_msa2)); - auto mlp_out = mlp->forward(ctx, modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp)); + auto mlp_out = mlp->forward(ctx, modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp)); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, gate_mlp)); return x; @@ -564,7 +563,7 @@ struct DismantledBlock : public GGMLBlock { attn_out = attn->post_attention(ctx, attn_out); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, attn_out, gate_msa)); - auto mlp_out = mlp->forward(ctx, modulate(ctx->ggml_ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp)); + auto mlp_out = mlp->forward(ctx, modulate(ctx, norm2->forward(ctx, x), shift_mlp, scale_mlp)); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, mlp_out, gate_mlp)); return x; @@ -755,7 +754,7 @@ struct FinalLayer : public GGMLBlock { auto shift = m_vec[0]; // [N, hidden_size] auto scale = m_vec[1]; // [N, hidden_size] - x = modulate(ctx->ggml_ctx, norm_final->forward(ctx, x), shift, scale); + x = modulate(ctx, norm_final->forward(ctx, x), shift, scale); x = linear->forward(ctx, x); return x; diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 955c5277a7..4ff91c416c 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -16,6 +16,11 @@ enum class RefImageResizeMode { AREA, }; +enum class RefImageSizeMode { + MATCH, + MAX, +}; + struct RefImageParams { bool pass_to_vlm = false; bool pass_to_dit = true; @@ -27,9 +32,26 @@ struct RefImageParams { int vlm_min_size = -1; int vlm_max_size = -1; bool resize_vae_to_target = false; + bool resize_single_ref_to_target = false; + bool vae_size_buckets = false; + bool vae_center_crop = false; + bool vlm_use_vae_input = false; + RefImageSizeMode ref_image_size = RefImageSizeMode::MATCH; }; const std::unordered_map REF_IMAGE_PRESETS = { + {"fibo2", [] { + RefImageParams params; + params.pass_to_vlm = true; + params.ref_index_mode = Rope::RefIndexMode::INCREASE; + params.vlm_resize_mode = RefImageResizeMode::LONGEST_SIDE; + params.vlm_max_size = 768; + params.resize_single_ref_to_target = true; + params.vae_size_buckets = true; + params.vae_center_crop = true; + params.vlm_use_vae_input = true; + return params; + }()}, {"flux_kontext", {false, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, {"longcat", {true, true, Rope::RefIndexMode::FIXED, false, true, -1, RefImageResizeMode::AREA, -1, -1}}, {"flux2", {false, true, Rope::RefIndexMode::INCREASE, false, true, -1, RefImageResizeMode::NONE, -1, -1}}, @@ -149,6 +171,11 @@ struct MingImageDiffusionExtra { const sd::Tensor* direct_context = nullptr; }; +struct ZImageDiffusionExtra { + const sd::Tensor* control_latent = nullptr; + float strength = 1.f; +}; + using DiffusionExtraParams = std::variant; + MingImageDiffusionExtra, + ZImageDiffusionExtra>; struct DiffusionParams { const sd::Tensor* x = nullptr; diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index f1cbbc7f37..a0fae7448f 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -2,8 +2,11 @@ #define __SD_MODEL_DIFFUSION_MOSS_IMAGE_HPP__ #include +#include #include +#include "core/pe_cache.h" + #include "model/common/ggml_block.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" @@ -53,13 +56,6 @@ namespace MOSSImage { } }; - inline ggml_tensor* situ(ggml_context* c, ggml_tensor* hidden, ggml_tensor* gate) { - auto a = ggml_scale(c, ggml_tanh(c, ggml_scale(c, gate, 0.25f)), 4.f); - a = ggml_mul(c, a, ggml_sigmoid(c, gate)); - auto b = ggml_scale(c, ggml_tanh(c, ggml_scale(c, hidden, 1.f / 25.f)), 25.f); - return ggml_mul(c, a, b); - } - struct TextBlock : public GGMLBlock { MOSSImageConfig config; @@ -83,7 +79,8 @@ namespace MOSSImage { auto y = norm("norm1", x); auto qkv = linear("qkv", y); auto split = [&](int i) { - return ggml_reshape_4d(c, ggml_ext_slice(c, qkv, 0, i * config.dim, (i + 1) * config.dim), config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2]); + // q/k feed rms_norm and v a permuted copy, which read these strided views in place. + return ggml_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2], config.head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * config.dim * qkv->nb[0]); }; auto q = norm("q_norm", split(0)); auto k = norm("k_norm", split(1)); @@ -92,7 +89,7 @@ namespace MOSSImage { y = linear("proj", y); x = ggml_add(c, x, y); y = norm("norm2", x); - y = situ(c, linear("ffn_up", y), linear("ffn_gate", y)); + y = ggml_ext_situ(ctx, linear("ffn_up", y), linear("ffn_gate", y)); y = linear("ffn_down", y); return ggml_add(c, x, y); } @@ -113,16 +110,17 @@ namespace MOSSImage { blocks["proj.ffn_down"] = std::make_shared(config.ffn_dim, config.dim, false); } - ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, ggml_tensor* mod) { + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, const std::array& mod) { auto c = ctx->ggml_ctx; auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; - auto part = [&](int i) { return ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim); }; + auto part = [&](int i) { return mod[i]; }; auto y = norm("norm1", x); - y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(1))), part(0)); + y = ggml_ext_modulate(ctx, y, part(0), part(1)); auto qkv = linear("proj.qkv", y); auto split = [&](int i) { - return ggml_reshape_4d(c, ggml_ext_slice(c, qkv, 0, i * config.dim, (i + 1) * config.dim), config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2]); + return ggml_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2], + config.head_dim * qkv->nb[0], qkv->nb[1], qkv->nb[2], i * config.dim * qkv->nb[0]); }; auto q = norm("q_norm", split(0)); auto k = norm("k_norm", split(1)); @@ -132,9 +130,9 @@ namespace MOSSImage { y = ggml_mul(c, y, part(2)); x = ggml_add(c, x, y); y = norm("norm2", x); - y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3)); + y = ggml_ext_modulate(ctx, y, part(3), part(4)); auto up_gate = linear("proj.ffn_up_gate", y); - y = situ(c, ggml_ext_slice(c, up_gate, 0, 0, config.ffn_dim), ggml_ext_slice(c, up_gate, 0, config.ffn_dim, 2 * config.ffn_dim)); + y = ggml_ext_situ(ctx, up_gate); y = linear("proj.ffn_down", y); y = ggml_mul(c, y, part(5)); return ggml_add(c, x, y); @@ -175,10 +173,17 @@ namespace MOSSImage { time = linear("time_embedding.2", ggml_silu(c, linear("time_embedding.0", time))); auto mod = linear("shared_block_modulation", time); auto head = linear("video_head_modulation", time); - x = ggml_concat(c, x, text, 1); + std::array parts; + for (int i = 0; i < 6; ++i) { + parts[i] = ggml_ext_slice(c, mod, 0, i * config.dim, (i + 1) * config.dim); + if (ctx->graph != nullptr) { + ggml_build_forward_expand(ctx->graph, parts[i]); + } + } + x = ggml_concat(c, x, text, 1); for (int64_t i = 0; i < config.depth; ++i) { if (std::find(skip.begin(), skip.end(), i) == skip.end()) { - x = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, x, pe, mod); + x = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, x, pe, parts); sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x"); } } @@ -186,16 +191,17 @@ namespace MOSSImage { x = std::dynamic_pointer_cast(blocks["final_norm_video"])->forward(ctx, x); auto shift = ggml_ext_slice(c, head, 0, 0, config.dim); auto scale = ggml_ext_slice(c, head, 0, config.dim, 2 * config.dim); - x = ggml_add(c, ggml_add(c, x, ggml_mul(c, x, scale)), shift); + x = ggml_ext_modulate(ctx, x, shift, scale); x = linear("final_linear_video", x); return DiT::unpatchify(c, x, image->ne[1], image->ne[0], 1, 1); } }; struct MOSSImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; MOSSImageConfig config; Model model; - std::vector pe_data, time_data; + std::vector time_data; MOSSImageRunner(ggml_backend_t backend, const String2TensorStorage& weights, const std::string& prefix, std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(MOSSImageConfig::detect_from_weights(weights, prefix)), model(config) { @@ -224,33 +230,36 @@ namespace MOSSImage { auto text = make_input(*inputs.context); GGML_ASSERT(x->ne[3] == 1 && x->ne[2] == config.channels && text->ne[0] == config.text_dim); int64_t h = x->ne[1], w = x->ne[0], nt = text->ne[1]; - Rope::Embedding embedding; - embedding.positions.append_image(static_cast(h), static_cast(w)); - embedding.positions.images.back().height_axis = 2; - embedding.positions.images.back().width_axis = 3; - embedding.positions.append_tokens(nt); - for (int64_t token = 0; token < h * w + nt; ++token) { - float pos[4] = {0.f, float(token - h * w), 0.f, 0.f}; - if (token < h * w) { - pos[1] = float(nt); - pos[2] = float(token / w) - float(h - 1) / 2.f; - pos[3] = float(token % w) - float(w - 1) / 2.f; + sd::PECacheKey pe_key("moss_image.pe"); + pe_key.add(h).add(w).add(nt).add(config.head_dim).add(circular_x_enabled).add(circular_y_enabled); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + Rope::Embedding embedding; + embedding.positions.append_image(static_cast(h), static_cast(w)); + embedding.positions.images.back().height_axis = 2; + embedding.positions.images.back().width_axis = 3; + embedding.positions.append_tokens(nt); + for (int64_t token = 0; token < h * w + nt; ++token) { + float pos[4] = {0.f, float(token - h * w), 0.f, 0.f}; + if (token < h * w) { + pos[1] = float(nt); + pos[2] = float(token / w) - float(h - 1) / 2.f; + pos[3] = float(token % w) - float(w - 1) / 2.f; + } + embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]}); } - embedding.ids.push_back({pos[0], pos[1], pos[2], pos[3]}); - } - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); - pe_data = finish_rope_pe(std::move(embedding)); + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, std::vector(4, static_cast(config.head_dim / 4)), embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + return sd::Tensor({2, 2, config.head_dim / 2, h * w + nt}, std::move(values)); + })); time_data.resize(config.freq_dim); for (int64_t i = 0; i < config.freq_dim / 2; ++i) { double angle = inputs.timesteps->values()[0] * std::pow(10000.0, -double(i) / double(config.freq_dim / 2)); time_data[i] = float(std::cos(angle)); time_data[i + config.freq_dim / 2] = float(std::sin(angle)); } - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, h * w + nt); auto time = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.freq_dim, 1, 1); - set_backend_tensor_data(pe, pe_data.data()); set_backend_tensor_data(time, time_data.data()); - auto ctx = get_context(); + auto ctx = get_context(gf); ggml_build_forward_expand(gf, model.forward(&ctx, x, time, text, pe, skip)); return gf; }; diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 6edc83422d..0e502371a8 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -7,6 +7,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -252,8 +253,8 @@ namespace Pid { auto w1 = std::dynamic_pointer_cast(blocks["w1"]); auto w2 = std::dynamic_pointer_cast(blocks["w2"]); auto w3 = std::dynamic_pointer_cast(blocks["w3"]); - auto h = ggml_silu_inplace(ctx->ggml_ctx, w1->forward(ctx, x)); - h = ggml_mul_inplace(ctx->ggml_ctx, h, w3->forward(ctx, x)); + auto h = w1->forward(ctx, x); + h = ggml_swiglu_split(ctx->ggml_ctx, h, w3->forward(ctx, x)); return w2->forward(ctx, h); } }; @@ -644,7 +645,7 @@ namespace Pid { int64_t effective_down_factor = config.lq_latent_down_factor / unpatchify_factor; float z_to_patch_ratio = static_cast(config.lq_sr_scale * effective_down_factor) / - static_cast(config.patch_size); + static_cast(config.patch_size); GGML_ASSERT(z_to_patch_ratio >= 1.0f); if (lq_latent->ne[0] != target_pW || lq_latent->ne[1] != target_pH) { lq_latent = ggml_interpolate(ctx->ggml_ctx, @@ -822,12 +823,9 @@ namespace Pid { }; struct PiDRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; PixelDiTConfig config; PixelDiT model; - std::vector pos_img_vec; - std::vector pos_txt_vec; - std::vector pixel_pos_vec; - std::vector pixel_pos_comp_vec; PiDRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map, @@ -867,57 +865,53 @@ namespace Pid { int64_t Hs = Hp / config.patch_size; int64_t Ws = Wp / config.patch_size; - pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.hidden_size / config.num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w))); - auto pos_img = ggml_new_tensor_4d(compute_ctx, - GGML_TYPE_F32, - 2, - 2, - config.hidden_size / config.num_groups / 2, - Hs * Ws); - set_backend_tensor_data(pos_img, pos_img_vec.data()); + sd::PECacheKey pos_img_key("pid.pos_img"); + pos_img_key.add(Hs).add(Ws).add(config.hidden_size).add(config.num_groups).add(config.rope_ref_grid_h); + pos_img_key.add(config.rope_ref_grid_w).add(circular_x_enabled).add(circular_y_enabled); + auto pos_img = make_input(pe_cache_.get_or_create(pos_img_key, [&]() { + auto pos_img_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.hidden_size / config.num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); + return sd::Tensor({2, 2, config.hidden_size / config.num_groups / 2, Hs * Ws}, std::move(pos_img_vec)); + })); int64_t Ltxt = std::min(context->ne[1], config.txt_max_length); - pos_txt_vec = make_rope_1d(static_cast(Ltxt), - static_cast(config.hidden_size / config.num_groups), - config.text_rope_theta); - auto pos_txt = ggml_new_tensor_4d(compute_ctx, - GGML_TYPE_F32, - 2, - 2, - config.hidden_size / config.num_groups / 2, - Ltxt); - set_backend_tensor_data(pos_txt, pos_txt_vec.data()); - - pixel_pos_vec = make_pixel_abs_pos(static_cast(Hp), - static_cast(Wp), - static_cast(config.pixel_hidden_size)); - auto pixel_pos = ggml_new_tensor_3d(compute_ctx, - GGML_TYPE_F32, - config.pixel_hidden_size, - Wp * Hp, - 1); - set_backend_tensor_data(pixel_pos, pixel_pos_vec.data()); - - pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), - static_cast(Ws), - static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), - 10000.f, - 16.f, - static_cast(config.rope_ref_grid_h), - static_cast(config.rope_ref_grid_w))); - auto pixel_pos_comp = ggml_new_tensor_4d(compute_ctx, - GGML_TYPE_F32, - 2, - 2, - config.pixel_attn_hidden_size / config.pixel_num_groups / 2, - Hs * Ws); - set_backend_tensor_data(pixel_pos_comp, pixel_pos_comp_vec.data()); + sd::PECacheKey pos_txt_key("pid.pos_txt"); + pos_txt_key.add(Ltxt).add(config.hidden_size).add(config.num_groups).add(config.text_rope_theta); + auto pos_txt = make_input(pe_cache_.get_or_create(pos_txt_key, [&]() { + auto pos_txt_vec = make_rope_1d(static_cast(Ltxt), + static_cast(config.hidden_size / config.num_groups), + config.text_rope_theta); + return sd::Tensor({2, 2, config.hidden_size / config.num_groups / 2, Ltxt}, std::move(pos_txt_vec)); + })); + + sd::PECacheKey pixel_pos_key("pid.pixel_pos"); + pixel_pos_key.add(Hp).add(Wp).add(config.pixel_hidden_size); + auto pixel_pos = make_input(pe_cache_.get_or_create(pixel_pos_key, [&]() { + auto pixel_pos_vec = make_pixel_abs_pos(static_cast(Hp), + static_cast(Wp), + static_cast(config.pixel_hidden_size)); + return sd::Tensor({config.pixel_hidden_size, Wp * Hp, 1}, std::move(pixel_pos_vec)); + })); + + sd::PECacheKey pixel_pos_comp_key("pid.pixel_pos_comp"); + pixel_pos_comp_key.add(Hs).add(Ws).add(config.pixel_attn_hidden_size).add(config.pixel_num_groups); + pixel_pos_comp_key.add(config.rope_ref_grid_h).add(config.rope_ref_grid_w).add(circular_x_enabled); + pixel_pos_comp_key.add(circular_y_enabled); + auto pixel_pos_comp = make_input(pe_cache_.get_or_create(pixel_pos_comp_key, [&]() { + auto pixel_pos_comp_vec = finish_rope_pe(make_rope_2d(static_cast(Hs), + static_cast(Ws), + static_cast(config.pixel_attn_hidden_size / config.pixel_num_groups), + 10000.f, + 16.f, + static_cast(config.rope_ref_grid_h), + static_cast(config.rope_ref_grid_w))); + return sd::Tensor({2, 2, config.pixel_attn_hidden_size / config.pixel_num_groups / 2, Hs * Ws}, std::move(pixel_pos_comp_vec)); + })); auto runner_ctx = get_context(); auto out = model.forward(&runner_ctx, diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp index cc06291928..dc6e2c2e7d 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -5,6 +5,7 @@ #include #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_runner.h" @@ -197,11 +198,11 @@ namespace PixArt { return ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, y, g)); }; - auto h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]); + auto h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[0], mv[1]); x = gate(attn1->forward(ctx, h, h), mv[2]); // ada_norm_single: no norm before cross-attention (PixArtMS.py) x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask)); - h = modulate(ctx->ggml_ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]); + h = modulate(ctx, norm(ctx->ggml_ctx, x), mv[3], mv[4]); h = proj->forward(ctx, h); h = ggml_ext_gelu(ctx->ggml_ctx, h, true); h = fc2->forward(ctx, h); @@ -224,9 +225,9 @@ namespace PixArt { PixArtModel(const PixArtConfig& config) : config(config) { blocks["x_embedder.proj"] = std::make_shared(config.in_channels, - config.hidden_size, - std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, - std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}); + config.hidden_size, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}, + std::pair{static_cast(config.patch_size), static_cast(config.patch_size)}); blocks["t_embedder"] = std::make_shared(ADALN_EMBED_DIM, config.hidden_size); blocks["t_block.1"] = std::make_shared(config.hidden_size, 6 * config.hidden_size); blocks["y_embedder.y_proj.fc1"] = std::make_shared(config.caption_channels, config.hidden_size); @@ -237,7 +238,7 @@ namespace PixArt { } blocks["final_layer.norm_final"] = std::make_shared(config.hidden_size, 1e-6f, false); blocks["final_layer.linear"] = std::make_shared(config.hidden_size, - config.patch_size * config.patch_size * config.out_channels); + config.patch_size * config.patch_size * config.out_channels); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -287,7 +288,7 @@ namespace PixArt { ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N), 2, 0); h = std::dynamic_pointer_cast(blocks["final_layer.norm_final"])->forward(ctx, h); - h = modulate(ctx->ggml_ctx, h, parts[0], parts[1]); + h = modulate(ctx, h, parts[0], parts[1]); h = std::dynamic_pointer_cast(blocks["final_layer.linear"])->forward(ctx, h); // [N, hp*wp, p*p*out_ch] h = DiT::unpatchify(ctx->ggml_ctx, h, hp, wp, static_cast(p), static_cast(p), false); return h; // [N, out_channels, H, W] @@ -295,9 +296,9 @@ namespace PixArt { }; struct PixArtRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; PixArtConfig config; PixArtModel model; - std::vector pos_vec; PixArtRunner(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, @@ -357,10 +358,13 @@ namespace PixArt { int64_t wp = W / config.patch_size; int64_t hp = H / config.patch_size; - pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp, - config.pos_embed_base_size, config.interpolation_scale); - auto pos = ggml_new_tensor_3d(compute_ctx, GGML_TYPE_F32, config.hidden_size, wp * hp, 1); - set_backend_tensor_data(pos, pos_vec.data()); + sd::PECacheKey pos_key("pixart.pos"); + pos_key.add(hp).add(wp).add(config.hidden_size).add(config.pos_embed_base_size).add(config.interpolation_scale); + auto pos = make_input(pe_cache_.get_or_create(pos_key, [&]() { + auto pos_vec = gen_2d_sincos_pos_embed(config.hidden_size, hp, wp, + config.pos_embed_base_size, config.interpolation_scale); + return sd::Tensor({config.hidden_size, wp * hp, 1}, std::move(pos_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask); diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 82d1730059..6435222f59 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -4,6 +4,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "core/util.h" #include "model/common/block.hpp" @@ -349,11 +350,11 @@ namespace Qwen { auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params); auto img_normed = img_norm1->forward(ctx, img); - auto img_modulated = Flux::modulate(ctx->ggml_ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr); + auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr); auto img_gate1 = img_mod_param_vec[2]; auto txt_normed = txt_norm1->forward(ctx, txt); - auto txt_modulated = Flux::modulate(ctx->ggml_ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]); + auto txt_modulated = Flux::modulate(ctx, txt_normed, txt_mod_param_vec[0], txt_mod_param_vec[1]); auto txt_gate1 = txt_mod_param_vec[2]; auto [img_attn_output, txt_attn_output] = attn->forward(ctx, img_modulated, txt_modulated, pe); @@ -362,11 +363,11 @@ namespace Qwen { txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, txt_attn_output, txt_gate1)); auto img_normed2 = img_norm2->forward(ctx, img); - auto img_modulated2 = Flux::modulate(ctx->ggml_ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr); + auto img_modulated2 = Flux::modulate(ctx, img_normed2, img_mod_param_vec[3], img_mod_param_vec[4], modulate_index != nullptr); auto img_gate2 = img_mod_param_vec[5]; auto txt_normed2 = txt_norm2->forward(ctx, txt); - auto txt_modulated2 = Flux::modulate(ctx->ggml_ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]); + auto txt_modulated2 = Flux::modulate(ctx, txt_normed2, txt_mod_param_vec[3], txt_mod_param_vec[4]); auto txt_gate2 = txt_mod_param_vec[5]; auto img_mlp_out = img_mlp->forward(ctx, img_modulated2); @@ -406,7 +407,7 @@ namespace Qwen { auto shift = mods[1]; x = norm->forward(ctx, x); - x = Flux::modulate(ctx->ggml_ctx, x, shift, scale); + x = Flux::modulate(ctx, x, shift, scale); return x; } @@ -561,10 +562,11 @@ namespace Qwen { }; struct QwenImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: QwenImageConfig config; QwenImageModel qwen_image; - std::vector pe_vec; std::vector modulate_index_vec; std::vector additional_t_cond_vec; SDVersion version; @@ -635,23 +637,24 @@ namespace Qwen { ref_index_mode = Rope::RefIndexMode::DECREASE; } - pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, - static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - batch_size, - static_cast(context->ne[1]), - ref_latents, - ref_index_mode, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe, true, "pe"); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("qwen_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(time_len).add(batch_size).add(config.patch_size); + pe_key.add_shapes(ref_latents).add(ref_index_mode); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_qwen_image_pe(time_len, + static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + batch_size, + static_cast(context->ne[1]), + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); ggml_tensor* modulate_index = nullptr; if (config.zero_cond_t) { diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 1d37d6afeb..76cbccb838 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -1,6 +1,7 @@ #ifndef __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__ #define __SD_MODEL_DIFFUSION_QWEN_IMAGE_2_1_H__ +#include "core/pe_cache.h" #include "model/diffusion/qwen_image.hpp" namespace Qwen { @@ -265,16 +266,18 @@ namespace Qwen { blocks["img_mlp.out"] = std::make_shared(config.intermediate_size, config.hidden_size, false); } - static ggml_tensor* modulate(ggml_context* ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) { + static ggml_tensor* modulate(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* params, int64_t prefix_length, bool gate = false) { + auto ctx = runner_ctx->ggml_ctx; auto rows = ggml_ext_chunk(ctx, params, 2, 1); auto apply = [&](ggml_tensor* part, ggml_tensor* row) { row = gate ? ggml_tanh(ctx, row) : ggml_scale_bias(ctx, row, 1.f, 1.f); + runner_ctx->expand_graph(row); return ggml_mul(ctx, part, row); }; - auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]); if (prefix_length == 0) { - return target; + return apply(x, rows[0]); } + auto target = apply(ggml_ext_slice(ctx, x, 1, prefix_length, x->ne[1]), rows[0]); auto prefix = apply(ggml_ext_slice(ctx, x, 1, 0, prefix_length), rows[1]); return ggml_concat(ctx, prefix, target, 1); } @@ -282,25 +285,22 @@ namespace Qwen { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, const std::vector& modulation, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { const int64_t prefix_length = cache.mode == QwenImage21PrefixCache::Mode::REUSE ? 0 : layout.prefix_length; auto h = std::dynamic_pointer_cast(blocks["img_norm1"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[0], prefix_length); + h = modulate(ctx, h, modulation[0], prefix_length); h = std::dynamic_pointer_cast(blocks["attn"])->forward(ctx, h, pe, layout.segments, masks, cache); - x = ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[1], prefix_length, true)); + x = ggml_add(ctx->ggml_ctx, x, modulate(ctx, h, modulation[1], prefix_length, true)); h = std::dynamic_pointer_cast(blocks["img_norm2"])->forward(ctx, x); - h = modulate(ctx->ggml_ctx, h, modulation[2], prefix_length); - ggml_tensor* gate; - auto fused = blocks.find("img_mlp.gate_up"); + h = modulate(ctx, h, modulation[2], prefix_length); + auto fused = blocks.find("img_mlp.gate_up"); if (fused != blocks.end()) { auto gate_up = std::dynamic_pointer_cast(fused->second)->forward(ctx, h); - auto parts = ggml_ext_chunk(ctx->ggml_ctx, gate_up, 2, 0); - gate = parts[0]; - h = parts[1]; + h = ggml_swiglu(ctx->ggml_ctx, gate_up); } else { - gate = std::dynamic_pointer_cast(blocks["img_mlp.gate_layer"])->forward(ctx, h); - h = std::dynamic_pointer_cast(blocks["img_mlp.proj"])->forward(ctx, h); + auto gate = std::dynamic_pointer_cast(blocks["img_mlp.gate_layer"])->forward(ctx, h); + h = std::dynamic_pointer_cast(blocks["img_mlp.proj"])->forward(ctx, h); + h = ggml_swiglu_split(ctx->ggml_ctx, gate, h); } - h = ggml_mul(ctx->ggml_ctx, h, ggml_silu(ctx->ggml_ctx, gate)); h = std::dynamic_pointer_cast(blocks["img_mlp.out"])->forward(ctx, h); - return ggml_add(ctx->ggml_ctx, x, modulate(ctx->ggml_ctx, h, modulation[3], prefix_length, true)); + return ggml_add(ctx->ggml_ctx, x, modulate(ctx, h, modulation[3], prefix_length, true)); } }; @@ -370,9 +370,9 @@ namespace Qwen { }; struct QwenImage21Runner : public DiffusionModelRunner { + sd::PECache pe_cache_; QwenImage21Config config; QwenImage21Model model; - std::vector pe_data; std::vector> mask_data; ggml_type prefix_cache_type = GGML_TYPE_COUNT; bool prefix_cache_enabled = true; @@ -486,8 +486,8 @@ namespace Qwen { } QwenImage21PrefixCache cache; if (prefix_cache_enabled && !prefix_cache_disabled && extra != nullptr && extra->prefix_id != 0 && layout.prefix_length > 0) { - cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) + - ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled); + cache.name = "qwen_image_2_1.prefix." + std::to_string(extra->prefix_id) + + ".circular." + std::to_string(circular_x_enabled) + std::to_string(circular_y_enabled); cache.prefix_length = layout.prefix_length; if (prefix_cache_type != GGML_TYPE_COUNT) { cache.type = prefix_cache_type; @@ -504,20 +504,7 @@ namespace Qwen { if (prefix_cache_type == GGML_TYPE_COUNT && active_cache.type == GGML_TYPE_F16) { checked_cache.flash_attn_used = &flash_attn_used; } - const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; - const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); - Rope::Embedding embedding; - embedding.ids.assign(first_position, layout.positions.end()); - const size_t offset = cached ? static_cast(layout.prefix_length) : 0; - embedding.positions.token_count = embedding.ids.size(); - for (auto region : layout.rope_layout.images) { - if (region.begin >= offset) { - region.begin -= offset; - embedding.positions.images.push_back(region); - } - } - embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies); - pe_data = finish_rope_pe(std::move(embedding)); + const bool cached = active_cache.mode == QwenImage21PrefixCache::Mode::REUSE; mask_data.clear(); if (!cached) { for (const auto& segment : layout.segments) { @@ -535,9 +522,30 @@ namespace Qwen { } auto build = [&]() { auto graph = new_graph_custom(QWEN_IMAGE_GRAPH_SIZE * 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.head_dim / 2, - layout.positions.size() - (cached ? layout.prefix_length : 0)); - set_backend_tensor_data(pe, pe_data.data()); + sd::PECacheKey pe_key("qwen_image_2_1.pe"); + pe_key.add(layout.positions).add(cached).add(layout.prefix_length).add(config.axes_dim).add(config.head_dim); + pe_key.add(circular_x_enabled).add(circular_y_enabled); + for (const auto& region : layout.rope_layout.images) { + pe_key.add(region.begin).add(region.count).add(region.height_period).add(region.width_period); + pe_key.add(region.height_axis).add(region.width_axis); + } + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + const auto first_position = layout.positions.begin() + (cached ? layout.prefix_length : 0); + Rope::Embedding embedding; + embedding.ids.assign(first_position, layout.positions.end()); + const size_t offset = cached ? static_cast(layout.prefix_length) : 0; + embedding.positions.token_count = embedding.ids.size(); + for (auto region : layout.rope_layout.images) { + if (region.begin >= offset) { + region.begin -= offset; + embedding.positions.images.push_back(region); + } + } + embedding.values = Rope::embed_nd(embedding.ids, 1, 10000.f, config.axes_dim, embedding.layout, &embedding.frequencies); + auto values = finish_rope_pe(std::move(embedding)); + const int64_t tokens = static_cast(layout.positions.size()) - (cached ? layout.prefix_length : 0); + return sd::Tensor({2, 2, config.head_dim / 2, tokens}, std::move(values)); + })); std::vector masks, ref_inputs; for (const auto& mask : mask_data) { masks.push_back(mask.empty() ? nullptr : make_input(mask)); @@ -549,7 +557,7 @@ namespace Qwen { } auto ctx = get_context(graph); auto out = model.forward(&ctx, make_input(x), make_input(*inputs.timesteps), cached ? nullptr : make_input(context), - ref_inputs, pe, layout, masks, checked_cache); + ref_inputs, pe, layout, masks, checked_cache); if (!flash_attn_used) { return static_cast(nullptr); } diff --git a/src/model/diffusion/sensenova_u1.h b/src/model/diffusion/sensenova_u1.h index 1702ed8546..302df6b26c 100644 --- a/src/model/diffusion/sensenova_u1.h +++ b/src/model/diffusion/sensenova_u1.h @@ -563,7 +563,7 @@ namespace SenseNovaU1 { blocks["language_model.model"] = std::make_shared(config); blocks["fm_modules.vision_model_mot_gen.embeddings"] = std::make_shared(config); blocks["fm_modules.timestep_embedder"] = std::make_shared(config.hidden_size, - config.timestep_embedding_size); + config.timestep_embedding_size); if (config.add_noise_scale_embedding) { blocks["fm_modules.noise_scale_embedder"] = std::make_shared(config.hidden_size, config.timestep_embedding_size); diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 4b75203d35..3c76c4ea19 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -8,6 +8,7 @@ #include #include "core/ggml_extend_backend.h" #include "core/ggml_tensor_utils.h" +#include "core/pe_cache.h" #include "model/common/block.hpp" #include "model/common/rope.hpp" @@ -885,11 +886,12 @@ namespace WAN { }; struct WanRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: std::string desc = "wan"; WanConfig config; Wan wan; - std::vector pe_vec; SDVersion version; WanRunner(ggml_backend_t backend, @@ -1000,37 +1002,36 @@ namespace WAN { ggml_tensor* audio_embed = make_optional_input(audio_embed_tensor); ggml_tensor* ref_latent = make_optional_input(ref_latent_tensor); - pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), - static_cast(x->ne[1]), - static_cast(x->ne[0]), - std::get<0>(config.patch_size), - std::get<1>(config.patch_size), - std::get<2>(config.patch_size), - 1, - config.theta, - config.axes_dim); - if (ref_latent != nullptr) { - // Match S2V's reference-frame temporal offset. - int t_start = std::max(30, static_cast(x->ne[2]) + 9); - auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), - static_cast(ref_latent->ne[1]), - static_cast(ref_latent->ne[0]), + sd::PECacheKey pe_key("wan.pe"); + pe_key.add_shape(x).add_shape(ref_latent).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(std::get<0>(config.patch_size)).add(std::get<1>(config.patch_size)).add(std::get<2>(config.patch_size)); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = Rope::gen_wan_pe(static_cast(x->ne[2]), + static_cast(x->ne[1]), + static_cast(x->ne[0]), std::get<0>(config.patch_size), std::get<1>(config.patch_size), std::get<2>(config.patch_size), 1, config.theta, - config.axes_dim, - t_start); - pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end()); - } - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + config.axes_dim); + if (ref_latent != nullptr) { + int t_start = std::max(30, static_cast(x->ne[2]) + 9); + auto ref_pe = Rope::gen_wan_pe(static_cast(ref_latent->ne[2]), + static_cast(ref_latent->ne[1]), + static_cast(ref_latent->ne[0]), + std::get<0>(config.patch_size), + std::get<1>(config.patch_size), + std::get<2>(config.patch_size), + 1, + config.theta, + config.axes_dim, + t_start); + pe_vec.insert(pe_vec.end(), ref_pe.begin(), ref_pe.end()); + } + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); if (c_concat != nullptr) { x = ggml_concat(compute_ctx, x, c_concat, 3); diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index ed2644047e..1c66c96b7f 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -3,6 +3,7 @@ #include #include +#include "core/pe_cache.h" #include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" @@ -23,9 +24,10 @@ #endif namespace ZImage { - constexpr int Z_IMAGE_GRAPH_SIZE = 20480; - constexpr int ADALN_EMBED_DIM = 256; - constexpr int SEQ_MULTI_OF = 32; + constexpr int Z_IMAGE_GRAPH_SIZE = 24576; + constexpr int ADALN_EMBED_DIM = 256; + constexpr int SEQ_MULTI_OF = 32; + constexpr int CONTROL_IN_CHANNELS = 33; struct ZImageConfig { int patch_size = 2; @@ -292,17 +294,16 @@ namespace ZImage { } }; - __STATIC_INLINE__ ggml_tensor* modulate(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* modulate(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale) { // x: [N, L, C] // scale: [N, C], or [N, L, C] when the caller modulates per token (LLaDA-Image editing // feeds a per-token timestep embedding so each segment carries its own modulation). if (scale->ne[1] != x->ne[1]) { - scale = ggml_reshape_3d(ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] + scale = ggml_reshape_3d(ctx->ggml_ctx, scale, scale->ne[0], 1, scale->ne[1]); // [N, 1, C] } - x = ggml_add(ctx, x, ggml_mul(ctx, x, scale)); - return x; + return ggml_ext_modulate(ctx, x, nullptr, scale); } struct JointTransformerBlock : public GGMLBlock { @@ -359,14 +360,14 @@ namespace ZImage { auto gate_mlp = mods[3]; auto residual = x; - x = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, x), scale_msa); + x = modulate(ctx, attention_norm1->forward(ctx, x), scale_msa); x = attention->forward(ctx, x, pe, mask); x = attention_norm2->forward(ctx, x); x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_msa)); x = ggml_add(ctx->ggml_ctx, x, residual); residual = x; - x = modulate(ctx->ggml_ctx, ffn_norm1->forward(ctx, x), scale_mlp); + x = modulate(ctx, ffn_norm1->forward(ctx, x), scale_mlp); x = feed_forward->forward(ctx, x); x = ffn_norm2->forward(ctx, x); x = ggml_mul(ctx->ggml_ctx, x, ggml_tanh(ctx->ggml_ctx, gate_mlp)); @@ -391,6 +392,135 @@ namespace ZImage { } }; + struct ZImageControlTransformerBlock : public GGMLBlock { + protected: + int block_id; + bool has_before_proj; + int64_t hidden_size; + int64_t head_dim; + int64_t num_heads; + int64_t num_kv_heads; + bool qk_norm; + + public: + ZImageControlTransformerBlock(int block_id, + int64_t hidden_size, + int64_t head_dim, + int64_t num_heads, + int64_t num_kv_heads, + int64_t multiple_of, + float ffn_dim_multiplier, + float norm_eps, + bool qk_norm, + bool with_before_proj = false, + bool norm_elementwise_affine = true, + float qk_norm_eps = 1e-6f) + : block_id(block_id), + has_before_proj(with_before_proj), + hidden_size(hidden_size), + head_dim(head_dim), + num_heads(num_heads), + num_kv_heads(num_kv_heads), + qk_norm(qk_norm) { + blocks["attention.to_q"] = std::make_shared(hidden_size, num_heads * head_dim, false); + blocks["attention.to_k"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["attention.to_v"] = std::make_shared(hidden_size, num_kv_heads * head_dim, false); + blocks["attention.to_out.0"] = std::make_shared(num_heads * head_dim, hidden_size, false); + if (qk_norm) { + blocks["attention.norm_q"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + blocks["attention.norm_k"] = std::make_shared(head_dim, qk_norm_eps, norm_elementwise_affine); + } + blocks["attention_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm1"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["attention_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["ffn_norm2"] = std::make_shared(hidden_size, norm_eps, norm_elementwise_affine); + blocks["feed_forward"] = std::make_shared(hidden_size, hidden_size, multiple_of, ffn_dim_multiplier); + blocks["adaLN_modulation.0"] = std::make_shared(MIN(hidden_size, ADALN_EMBED_DIM), 4 * hidden_size); + if (with_before_proj) { + blocks["before_proj"] = std::make_shared(hidden_size, hidden_size); + } + blocks["after_proj"] = std::make_shared(hidden_size, hidden_size); + } + + ggml_tensor* forward_attention(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* pe) { + auto to_q = std::dynamic_pointer_cast(blocks["attention.to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["attention.to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["attention.to_v"]); + auto out_proj = std::dynamic_pointer_cast(blocks["attention.to_out.0"]); + + if (out_proj->has_quantized_weight()) { + out_proj->set_scale(1.f / 16.f); + } + + int64_t n_token = c->ne[1]; + int64_t N = c->ne[2]; + + if (sd_backend_is(ctx->backend, "ROCm")) { + out_proj->set_force_prec_f32(true); + to_q->set_force_prec_f32(true); + to_k->set_force_prec_f32(true); + to_v->set_force_prec_f32(true); + } + + auto q = ggml_reshape_4d(ctx->ggml_ctx, to_q->forward(ctx, c), head_dim, num_heads, n_token, N); + auto k = ggml_reshape_4d(ctx->ggml_ctx, to_k->forward(ctx, c), head_dim, num_kv_heads, n_token, N); + auto v = ggml_reshape_4d(ctx->ggml_ctx, to_v->forward(ctx, c), head_dim, num_kv_heads, n_token, N); + + if (qk_norm) { + q = std::dynamic_pointer_cast(blocks["attention.norm_q"])->forward(ctx, q); + k = std::dynamic_pointer_cast(blocks["attention.norm_k"])->forward(ctx, k); + } + + auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.f / 128.f); + return out_proj->forward(ctx, out); + } + + std::pair forward(GGMLRunnerContext* ctx, + ggml_tensor* c, + ggml_tensor* pe, + ggml_tensor* skip, + ggml_tensor* t_emb) { + auto attention_norm1 = std::dynamic_pointer_cast(blocks["attention_norm1"]); + auto ffn_norm1 = std::dynamic_pointer_cast(blocks["ffn_norm1"]); + auto attention_norm2 = std::dynamic_pointer_cast(blocks["attention_norm2"]); + auto ffn_norm2 = std::dynamic_pointer_cast(blocks["ffn_norm2"]); + auto feed_forward = std::dynamic_pointer_cast(blocks["feed_forward"]); + auto adaLN_modulation_0 = std::dynamic_pointer_cast(blocks["adaLN_modulation.0"]); + auto after_proj = std::dynamic_pointer_cast(blocks["after_proj"]); + + if (block_id == 0 && has_before_proj) { + GGML_ASSERT(skip != nullptr); + auto before_proj = std::dynamic_pointer_cast(blocks["before_proj"]); + c = ggml_add(ctx->ggml_ctx, before_proj->forward(ctx, c), skip); + } + + auto m = adaLN_modulation_0->forward(ctx, t_emb); + auto mods = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0); + auto scale_msa = mods[0]; + auto gate_msa = mods[1]; + auto scale_mlp = mods[2]; + auto gate_mlp = mods[3]; + + auto residual = c; + c = modulate(ctx, attention_norm1->forward(ctx, c), scale_msa); + c = forward_attention(ctx, c, pe); + c = attention_norm2->forward(ctx, c); + c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_msa)); + c = ggml_add(ctx->ggml_ctx, c, residual); + + residual = c; + c = modulate(ctx, ffn_norm1->forward(ctx, c), scale_mlp); + c = feed_forward->forward(ctx, c); + c = ffn_norm2->forward(ctx, c); + c = ggml_mul(ctx->ggml_ctx, c, ggml_tanh(ctx->ggml_ctx, gate_mlp)); + c = ggml_add(ctx->ggml_ctx, c, residual); + + return {after_proj->forward(ctx, c), c}; + } + }; + struct FinalLayer : public GGMLBlock { public: FinalLayer(int64_t hidden_size, @@ -413,7 +543,7 @@ namespace ZImage { auto scale = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, hidden_size] x = norm_final->forward(ctx, x); - x = modulate(ctx->ggml_ctx, x, scale); + x = modulate(ctx, x, scale); x = linear->forward(ctx, x); return x; @@ -423,15 +553,24 @@ namespace ZImage { class ZImageModel : public GGMLBlock { protected: ZImageConfig config; + GGMLBlockMap control_blocks; + int control_num_layers = 0; + int control_num_refiner_layers = 0; void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { params["cap_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); params["x_pad_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.hidden_size); + + for (auto& [name, block] : control_blocks) { + block->init(ctx, tensor_storage_map, name); + } } public: + bool control_enabled = false; + ZImageModel() = default; - ZImageModel(ZImageConfig config) + ZImageModel(ZImageConfig config, const String2TensorStorage& tensor_storage_map = {}) : config(config) { blocks["x_embedder"] = std::make_shared(config.patch_size * config.patch_size * config.in_channels, config.hidden_size); blocks["t_embedder"] = std::make_shared(MIN(config.hidden_size, 1024), 256, 256); @@ -484,13 +623,154 @@ namespace ZImage { } blocks["final_layer"] = std::make_shared(config.hidden_size, config.patch_size, config.out_channels); + + if (tensor_storage_map.find("control_all_x_embedder.2-1.weight") != tensor_storage_map.end()) { + control_enabled = true; + + for (const auto& pair : tensor_storage_map) { + if (starts_with(pair.first, "control_layers.")) { + auto items = split_string(pair.first, '.'); + if (items.size() > 1) { + control_num_layers = std::max(control_num_layers, atoi(items[1].c_str()) + 1); + } + } else if (starts_with(pair.first, "control_noise_refiner.")) { + auto items = split_string(pair.first, '.'); + if (items.size() > 1) { + control_num_refiner_layers = std::max(control_num_refiner_layers, atoi(items[1].c_str()) + 1); + } + } + } + + GGML_ASSERT(control_num_layers > 0); + GGML_ASSERT(control_num_refiner_layers == config.num_refiner_layers); + + control_blocks["control_all_x_embedder.2-1"] = std::make_shared(config.patch_size * config.patch_size * CONTROL_IN_CHANNELS, config.hidden_size); + + for (int i = 0; i < control_num_refiner_layers; i++) { + std::string name = "control_noise_refiner." + std::to_string(i); + auto block = std::make_shared(i, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, + tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end()); + + control_blocks[name] = block; + } + + for (int i = 0; i < control_num_layers; i++) { + std::string name = "control_layers." + std::to_string(i); + auto block = std::make_shared(i, + config.hidden_size, + config.head_dim, + config.num_heads, + config.num_kv_heads, + config.multiple_of, + config.ffn_dim_multiplier, + config.norm_eps, + config.qk_norm, + tensor_storage_map.find(name + ".before_proj.weight") != tensor_storage_map.end()); + + control_blocks[name] = block; + } + } + } + + void get_control_param_tensors(std::map& tensors) { + for (auto& [name, block] : control_blocks) { + block->get_param_tensors(tensors, name); + } + } + + ggml_tensor* forward_control_refiner(GGMLRunnerContext* ctx, + ggml_tensor* control, + ggml_tensor* img, + ggml_tensor* img_pe, + ggml_tensor* t_emb, + float strength, + std::vector& ref_hints) { + auto all_x_embedder = std::dynamic_pointer_cast(control_blocks["control_all_x_embedder.2-1"]); + + auto c = all_x_embedder->forward(ctx, control); + sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.prelude", "control"); + + int64_t N = c->ne[2]; + int64_t n_token = c->ne[1]; + + int64_t n_pad_token = Rope::bound_mod(static_cast(n_token), SEQ_MULTI_OF); + if (n_pad_token > 0) { + auto img_pad_token = params["x_pad_token"]; + auto ctrl_pad_tokens = ggml_repeat_4d(ctx->ggml_ctx, img_pad_token, img_pad_token->ne[0], n_pad_token, N, 1); + c = ggml_concat(ctx->ggml_ctx, c, ctrl_pad_tokens, 1); + } + + for (int i = 0; i < control_num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(control_blocks["control_noise_refiner." + std::to_string(i)]); + + auto out = block->forward(ctx, c, img_pe, img, t_emb); + c = out.second; + ref_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength)); + sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.noise_refiner." + std::to_string(i), "control"); + } + + return c; + } + + ggml_tensor* forward_noise_refiner_hinted(GGMLRunnerContext* ctx, + ggml_tensor* img, + ggml_tensor* img_pe, + ggml_tensor* t_emb, + const std::vector& ref_hints) { + GGML_ASSERT(ref_hints.size() == static_cast(config.num_refiner_layers)); + + auto c = img; + for (int i = 0; i < config.num_refiner_layers; i++) { + auto block = std::dynamic_pointer_cast(blocks["noise_refiner." + std::to_string(i)]); + + c = block->forward(ctx, c, img_pe, nullptr, t_emb); + c = ggml_add(ctx->ggml_ctx, c, ref_hints[i]); + sd::ggml_graph_cut::mark_graph_cut(c, "z_image.noise_refiner_hinted." + std::to_string(i), "img"); + } + + return c; + } + + std::vector forward_control_main(GGMLRunnerContext* ctx, + ggml_tensor* control_ctx, + ggml_tensor* img_c, + ggml_tensor* txt, + ggml_tensor* pe, + ggml_tensor* t_emb, + float strength) { + auto unified_c = ggml_concat(ctx->ggml_ctx, txt, img_c, 1); + auto ctrl_unified = ggml_concat(ctx->ggml_ctx, txt, control_ctx, 1); + + std::vector main_hints; + main_hints.reserve(control_num_layers); + + for (int i = 0; i < control_num_layers; i++) { + auto block = std::dynamic_pointer_cast(control_blocks["control_layers." + std::to_string(i)]); + + auto out = block->forward(ctx, ctrl_unified, pe, unified_c, t_emb); + ctrl_unified = out.second; + main_hints.push_back(ggml_ext_scale(ctx->ggml_ctx, out.first, strength)); + sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control"); + } + + return main_hints; } ggml_tensor* forward_core(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* control = nullptr, + float strength = 0.f) { auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); auto t_embedder = std::dynamic_pointer_cast(blocks["t_embedder"]); auto cap_embedder_0 = std::dynamic_pointer_cast(blocks["cap_embedder.0"]); @@ -530,6 +810,14 @@ namespace ZImage { auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt->ne[1]); auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt->ne[1], pe->ne[3]); + std::vector ref_hints; + ggml_tensor* control_ctx = nullptr; + if (control != nullptr) { + control_ctx = forward_control_refiner(ctx, control, img, img_pe, t_emb, strength, ref_hints); + } + + ggml_tensor* img_pre_refiner = img; + for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); @@ -544,13 +832,27 @@ namespace ZImage { sd::ggml_graph_cut::mark_graph_cut(img, "z_image.noise_refiner." + std::to_string(i), "img"); } + ggml_tensor* img_c = nullptr; + if (control_ctx != nullptr) { + img_c = forward_noise_refiner_hinted(ctx, img_pre_refiner, img_pe, t_emb, ref_hints); + } + auto txt_img = ggml_concat(ctx->ggml_ctx, txt, img, 1); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, hidden_size] sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.prelude", "txt_img"); + std::vector main_hints; + if (control_ctx != nullptr) { + main_hints = forward_control_main(ctx, control_ctx, img_c, txt, pe, t_emb, strength); + } + GGML_ASSERT(main_hints.empty() || main_hints.size() == static_cast(control_num_layers)); + for (int i = 0; i < config.num_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)]); txt_img = block->forward(ctx, txt_img, pe, nullptr, t_emb); + if (i % 2 == 0 && (i / 2) < static_cast(main_hints.size())) { + txt_img = ggml_add(ctx->ggml_ctx, txt_img, main_hints[i / 2]); + } sd::ggml_graph_cut::mark_graph_cut(txt_img, "z_image.layers." + std::to_string(i), "txt_img"); } @@ -566,7 +868,9 @@ namespace ZImage { ggml_tensor* timestep, ggml_tensor* context, ggml_tensor* pe, - std::vector ref_latents = {}) { + std::vector ref_latents = {}, + ggml_tensor* control_latent = nullptr, + float control_strength = 0.f) { // Forward pass of DiT. // x: [N, C, H, W] // timestep: [N,] @@ -591,7 +895,16 @@ namespace ZImage { } } - auto out = forward_core(ctx, img, timestep, context, pe); + ggml_tensor* control = nullptr; + if (control_latent != nullptr) { + GGML_ASSERT(control_enabled); + GGML_ASSERT(control_latent->ne[2] == CONTROL_IN_CHANNELS); + GGML_ASSERT(control_latent->ne[0] == W && control_latent->ne[1] == H && control_latent->ne[3] == N); + GGML_ASSERT(ref_latents.empty()); + control = DiT::pad_and_patchify(ctx, control_latent, patch_size, patch_size, false); + } + + auto out = forward_core(ctx, img, timestep, context, pe, control, control_strength); out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, n_img_token); // [N, n_img_token, ph*pw*C] out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); // [N, C, H, W] @@ -603,10 +916,11 @@ namespace ZImage { }; struct ZImageRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; + public: ZImageConfig config; ZImageModel z_image; - std::vector pe_vec; std::vector timestep_vec; SDVersion version; @@ -617,7 +931,7 @@ namespace ZImage { std::shared_ptr weight_manager = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(ZImageConfig::detect_from_weights(tensor_storage_map, prefix)) { - z_image = ZImageModel(config); + z_image = ZImageModel(config, tensor_storage_map); z_image.init(params_ctx, tensor_storage_map, prefix); } @@ -629,11 +943,21 @@ namespace ZImage { z_image.get_param_tensors(tensors, prefix); } + bool control_enabled() const { + return z_image.control_enabled; + } + + void get_control_param_tensors(std::map& tensors) { + z_image.get_control_param_tensors(tensors); + } + ggml_cgraph* build_graph(const sd::Tensor& x_tensor, const sd::Tensor& timesteps_tensor, const sd::Tensor& context_tensor, const std::vector>& ref_latents_tensor = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& control_latent_tensor = {}, + float control_strength = 0.f) { ggml_cgraph* gf = new_graph_custom(Z_IMAGE_GRAPH_SIZE); ggml_tensor* x = make_input(x_tensor); ggml_tensor* timesteps = make_input(timesteps_tensor); @@ -646,23 +970,26 @@ namespace ZImage { ref_latents.push_back(make_input(ref_latent_tensor)); } - pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - SEQ_MULTI_OF, - ref_latents, - ref_index_mode, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - // LOG_VERBOSE("pos_len %d", pos_len); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - // pe->data = pe_vec.data(); - // print_ggml_tensor(pe, true, "pe"); - // pe->data = nullptr; - set_backend_tensor_data(pe, pe_vec.data()); + ggml_tensor* control_latent = make_optional_input(control_latent_tensor); + + sd::PECacheKey pe_key("z_image.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size).add_shapes(ref_latents); + pe_key.add(ref_index_mode); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + SEQ_MULTI_OF, + ref_latents, + ref_index_mode, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto runner_ctx = get_context(); ggml_tensor* out = z_image.forward(&runner_ctx, @@ -670,7 +997,9 @@ namespace ZImage { timesteps, context, pe, - ref_latents); + ref_latents, + control_latent, + control_strength); ggml_build_forward_expand(gf, out); @@ -682,12 +1011,14 @@ namespace ZImage { const sd::Tensor& timesteps, const sd::Tensor& context, const std::vector>& ref_latents = {}, - Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED) { + Rope::RefIndexMode ref_index_mode = Rope::RefIndexMode::FIXED, + const sd::Tensor& control_latent = {}, + float control_strength = 0.f) { // x: [N, in_channels, h, w] // timesteps: [N, ] // context: [N, max_position, hidden_size] auto get_graph = [&]() -> ggml_cgraph* { - return build_graph(x, timesteps, context, ref_latents, ref_index_mode); + return build_graph(x, timesteps, context, ref_latents, ref_index_mode, control_latent, control_strength); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); @@ -698,12 +1029,21 @@ namespace ZImage { GGML_ASSERT(diffusion_params.x != nullptr); GGML_ASSERT(diffusion_params.timesteps != nullptr); static const std::vector> empty_ref_latents; + static const sd::Tensor empty_control_latent; + + const auto* z_image_extra = std::get_if(&diffusion_params.extra); + const sd::Tensor& control_latent = (z_image_extra != nullptr && z_image_extra->control_latent != nullptr) ? *z_image_extra->control_latent + : empty_control_latent; + float control_strength = z_image_extra != nullptr ? z_image_extra->strength : 0.f; + return compute(n_threads, *diffusion_params.x, *diffusion_params.timesteps, tensor_or_empty(diffusion_params.context), diffusion_params.ref_latents && diffusion_params.ref_image_params.pass_to_dit ? *diffusion_params.ref_latents : empty_ref_latents, - diffusion_params.ref_image_params.ref_index_mode); + diffusion_params.ref_image_params.ref_index_mode, + control_latent, + control_strength); } void test() { diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp index 0eb0b14d3d..f44f2ff9f8 100644 --- a/src/model/diffusion/z_image_l2p.hpp +++ b/src/model/diffusion/z_image_l2p.hpp @@ -2,6 +2,7 @@ #define __SD_MODEL_DIFFUSION_Z_IMAGE_L2P_HPP__ #include +#include "core/pe_cache.h" #include "z_image.hpp" @@ -183,9 +184,9 @@ namespace ZImageL2P { }; struct ZImageL2PRunner : public DiffusionModelRunner { + sd::PECache pe_cache_; ZImageL2PConfig config; ZImageL2PModel model; - std::vector pe_vec; ZImageL2PRunner(ggml_backend_t backend, const String2TensorStorage& tensors, @@ -221,19 +222,23 @@ namespace ZImageL2P { auto context = make_input(*inputs.context); GGML_ASSERT(x->ne[3] == 1); GGML_ASSERT(x->ne[0] % config.patch_size == 0 && x->ne[1] % config.patch_size == 0); - pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), - static_cast(x->ne[0]), - config.patch_size, - static_cast(x->ne[3]), - static_cast(context->ne[1]), - ZImage::SEQ_MULTI_OF, - {}, - Rope::RefIndexMode::FIXED, - config.theta, - config.axes_dim)); - int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); - auto pe = ggml_new_tensor_4d(compute_ctx, GGML_TYPE_F32, 2, 2, config.axes_dim_sum / 2, pos_len); - set_backend_tensor_data(pe, pe_vec.data()); + sd::PECacheKey pe_key("z_image_l2p.pe"); + pe_key.add_shape(x).add_shape(context).add(config.theta).add(config.axes_dim).add(config.axes_dim_sum); + pe_key.add(circular_x_enabled).add(circular_y_enabled).add(config.patch_size); + auto pe = make_input(pe_cache_.get_or_create(pe_key, [&]() { + auto pe_vec = finish_rope_pe(Rope::gen_z_image_pe(static_cast(x->ne[1]), + static_cast(x->ne[0]), + config.patch_size, + static_cast(x->ne[3]), + static_cast(context->ne[1]), + ZImage::SEQ_MULTI_OF, + {}, + Rope::RefIndexMode::FIXED, + config.theta, + config.axes_dim)); + int pos_len = static_cast(pe_vec.size() / config.axes_dim_sum / 2); + return sd::Tensor({2, 2, config.axes_dim_sum / 2, pos_len}, std::move(pe_vec)); + })); auto ctx = get_context(); auto out = model.forward(&ctx, x, t, context, pe); ggml_build_forward_expand(gf, out); diff --git a/src/model/integrations/anima.cpp b/src/model/integrations/anima.cpp new file mode 100644 index 0000000000..2f9b2babc0 --- /dev/null +++ b/src/model/integrations/anima.cpp @@ -0,0 +1,43 @@ +#include "model/diffusion/anima.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_anima(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& anima_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_ANIMA; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "cosmos_reference"; + definition.name = "Anima"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_anima; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/boogu_image.cpp b/src/model/integrations/boogu_image.cpp new file mode 100644 index 0000000000..3d6e8c06c1 --- /dev/null +++ b/src/model/integrations/boogu_image.cpp @@ -0,0 +1,46 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/boogu.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_boogu_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} + +const std::vector& boogu_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_BOOGU_IMAGE; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "z_image_omni"; + definition.name = "Boogu Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLOW_PRED, 3.16f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_boogu_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ernie_image.cpp b/src/model/integrations/ernie_image.cpp new file mode 100644 index 0000000000..f1d185cddd --- /dev/null +++ b/src/model/integrations/ernie_image.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/ernie_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ernie_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& ernie_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_ERNIE_IMAGE; + definition.geometry = {1, 128, 3}; + definition.name = "Ernie Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 4.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_ernie_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/fibo2.cpp b/src/model/integrations/fibo2.cpp new file mode 100644 index 0000000000..d16a0cf851 --- /dev/null +++ b/src/model/integrations/fibo2.cpp @@ -0,0 +1,39 @@ +#include "model/diffusion/fibo2.hpp" +#include "conditioning/conditioner.hpp" +#include "conditioning/fibo2.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_fibo2(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& fibo2_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FIBO2; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "fibo2"; + definition.name = "Fibo-2.0"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_fibo2; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/flux.cpp b/src/model/integrations/flux.cpp new file mode 100644 index 0000000000..fc96da7a38 --- /dev/null +++ b/src/model/integrations/flux.cpp @@ -0,0 +1,167 @@ +#include "model/diffusion/flux.hpp" +#include "conditioning/conditioner.hpp" +#include "model/integrations/flux.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_flux(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + if (version == VERSION_CHROMA || version == VERSION_CHROMA_RADIANCE) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + false, + 1, + false, + weight_manager, + sd_ctx_params->model_args); + } else if (version == VERSION_OVIS_IMAGE) { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + } else { + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + } + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& flux_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FLUX; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flux"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLUX_FILL; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flux Fill"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.is_inpaint = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLUX_CONTROLS; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flux Control"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.is_control = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLEX_2; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Flex.2"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.is_inpaint = true; + definition.traits.is_control = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_CHROMA_RADIANCE; + definition.geometry = {1, 3, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Chroma Radiance"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.traits.fake_vae = true; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_CHROMA; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Chroma"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_OVIS_IMAGE; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "flux_kontext"; + definition.name = "Ovis Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_flux; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} + +float flux_guidance_flow_shift(const String2TensorStorage& weights) { + for (const auto& [name, tensor] : weights) { + if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) + return 1.15f; + } + return 1.f; +} diff --git a/src/model/integrations/flux.hpp b/src/model/integrations/flux.hpp new file mode 100644 index 0000000000..3b11888b08 --- /dev/null +++ b/src/model/integrations/flux.hpp @@ -0,0 +1,5 @@ +#ifndef __SD_MODEL_INTEGRATIONS_FLUX_HPP__ +#define __SD_MODEL_INTEGRATIONS_FLUX_HPP__ +#include "model.h" +float flux_guidance_flow_shift(const String2TensorStorage& weights); +#endif diff --git a/src/model/integrations/flux2.cpp b/src/model/integrations/flux2.cpp new file mode 100644 index 0000000000..d895df2e3c --- /dev/null +++ b/src/model/integrations/flux2.cpp @@ -0,0 +1,79 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/flux.hpp" +#include "model/integrations/flux.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_flux2(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool is_chroma = false; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& flux2_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FLUX2; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "Flux.2"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX2_SCHEDULER; + definition.build = build_flux2; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_FLUX2_KLEIN; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "Flux.2 klein"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX_FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX2_SCHEDULER; + definition.build = build_flux2; + definition.sampling.flow_shift_from_weights = flux_guidance_flow_shift; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SEFI_IMAGE; + definition.geometry = {1, 144, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "SeFi-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {SEFI_FLOW_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_flux2; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/flux3.cpp b/src/model/integrations/flux3.cpp new file mode 100644 index 0000000000..9ee8d3db49 --- /dev/null +++ b/src/model/integrations/flux3.cpp @@ -0,0 +1,38 @@ +#include "model/diffusion/flux3.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_flux3(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, version, "", false, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& flux3_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_FLUX3; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "flux2"; + definition.name = "Flux.3 Image Klein"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX3_FLOW_PRED, 6.93f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_flux3; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/hidream_o1.cpp b/src/model/integrations/hidream_o1.cpp new file mode 100644 index 0000000000..66f6a9fc3a --- /dev/null +++ b/src/model/integrations/hidream_o1.cpp @@ -0,0 +1,41 @@ +#include "model/diffusion/hidream_o1.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_hidream_o1(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model", + weight_manager); + return true; +} + +const std::vector& hidream_o1_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_HIDREAM_O1; + definition.geometry = {1, 3, 3}; + definition.name = "HiDream O1"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_hidream_o1; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/hunyuan_video.cpp b/src/model/integrations/hunyuan_video.cpp new file mode 100644 index 0000000000..85bb741f88 --- /dev/null +++ b/src/model/integrations/hunyuan_video.cpp @@ -0,0 +1,47 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/hunyuan.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_hunyuan_video(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} + +const std::vector& hunyuan_video_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_HUNYUAN_VIDEO; + definition.geometry = {1, 32, 3}; + definition.name = "Hunyuan Video"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::HUNYUAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 7.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_hunyuan_video; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ideogram4.cpp b/src/model/integrations/ideogram4.cpp new file mode 100644 index 0000000000..77baef2724 --- /dev/null +++ b/src/model/integrations/ideogram4.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/ideogram4.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ideogram4(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& ideogram4_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_IDEOGRAM4; + definition.geometry = {1, 128, 3}; + definition.name = "Ideogram 4"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = LOGIT_NORMAL_SCHEDULER; + definition.build = build_ideogram4; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/iris.cpp b/src/model/integrations/iris.cpp new file mode 100644 index 0000000000..a75c88ca4f --- /dev/null +++ b/src/model/integrations/iris.cpp @@ -0,0 +1,37 @@ +#include "model/diffusion/iris.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_iris(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, version, "", false, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& iris_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_IRIS; + definition.geometry = {16, 3, 3}; + definition.name = "Iris-3B"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {FLOW_PRED, 4.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_iris; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/krea2.cpp b/src/model/integrations/krea2.cpp new file mode 100644 index 0000000000..367044537a --- /dev/null +++ b/src/model/integrations/krea2.cpp @@ -0,0 +1,46 @@ +#include "model/diffusion/krea2.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_krea2(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& krea2_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_KREA2; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "krea2_ostris_edit"; + definition.name = "Krea2"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLUX_FLOW_PRED, 1.15f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_krea2; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/lens.cpp b/src/model/integrations/lens.cpp new file mode 100644 index 0000000000..da66d2cd8c --- /dev/null +++ b/src/model/integrations/lens.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/lens.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_lens(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& lens_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LENS; + definition.geometry = {1, 128, 3}; + definition.name = "Lens"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLUX_FLOW_PRED, 1.83f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_lens; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/lingbot_video.cpp b/src/model/integrations/lingbot_video.cpp new file mode 100644 index 0000000000..d4cc6cc557 --- /dev/null +++ b/src/model/integrations/lingbot_video.cpp @@ -0,0 +1,54 @@ +#include "model/diffusion/lingbot_video.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_lingbot_video(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool enable_vision = false; + for (const auto& [name, _] : tensor_storage_map) { + if (starts_with(name, "text_encoders.llm.visual.")) { + enable_vision = true; + break; + } + } + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& lingbot_video_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LINGBOT_VIDEO; + definition.geometry = {2, 16, 3}; + definition.name = "LingBot Video"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_lingbot_video; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/llada_image.cpp b/src/model/integrations/llada_image.cpp new file mode 100644 index 0000000000..58c47e8f87 --- /dev/null +++ b/src/model/integrations/llada_image.cpp @@ -0,0 +1,50 @@ +#include "model/diffusion/llada_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_llada_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "queryformer", + "text_projection", + "sigvq", + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& llada_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LLADA_IMAGE; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "llada_image"; + definition.name = "LLaDA-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 1.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = LLADA_IMAGE_SCHEDULER; + definition.build = build_llada_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} + +// The attention projections keep their diffusers names (JointAttention's split_qkv mode), so +// only the patch-size-keyed dicts need flattening. Latents arrive already patchified from the +// Flux2 VAE, so the only patch key is 1-1. diff --git a/src/model/integrations/longcat.cpp b/src/model/integrations/longcat.cpp new file mode 100644 index 0000000000..d371710c46 --- /dev/null +++ b/src/model/integrations/longcat.cpp @@ -0,0 +1,47 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/flux.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_longcat(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& longcat_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LONGCAT; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "longcat"; + definition.name = "Longcat-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLUX_FLOW_PRED, 3.0f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_longcat; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ltxav.cpp b/src/model/integrations/ltxav.cpp new file mode 100644 index 0000000000..9995239249 --- /dev/null +++ b/src/model/integrations/ltxav.cpp @@ -0,0 +1,45 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/ltxv.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ltxav(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + "text_encoders.llm", + "text_embedding_projection", + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& ltxav_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_LTXAV; + definition.geometry = {1, 128, 3}; + definition.name = "LTXAV"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::LTXAV; + definition.traits.video_tae = true; + definition.sampling = {FLUX_FLOW_PRED, 2.37f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = LTX2_SCHEDULER; + definition.build = build_ltxav; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/mage_flow.cpp b/src/model/integrations/mage_flow.cpp new file mode 100644 index 0000000000..000fa1571b --- /dev/null +++ b/src/model/integrations/mage_flow.cpp @@ -0,0 +1,45 @@ +#include "model/diffusion/mage_flow.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_mage_flow(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& mage_flow_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MAGE_FLOW; + definition.geometry = {1, 128, 3}; + definition.ref_image_preset = "mage_flow"; + definition.name = "Mage Flow"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::MAGE; + definition.sampling = {FLOW_PRED, 6.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_mage_flow; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/ming_image.cpp b/src/model/integrations/ming_image.cpp new file mode 100644 index 0000000000..8587df8c8a --- /dev/null +++ b/src/model/integrations/ming_image.cpp @@ -0,0 +1,38 @@ +#include "model/diffusion/ming_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_ming_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, weight_manager, tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, "model.diffusion_model", weight_manager); + return true; +} + +const std::vector& ming_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MING_IMAGE; + definition.geometry = {2, 16, 4}; + definition.name = "Ming-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {MING_IMAGE_FLOW_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_ming_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/minimax_h3.cpp b/src/model/integrations/minimax_h3.cpp new file mode 100644 index 0000000000..35a1f6700b --- /dev/null +++ b/src/model/integrations/minimax_h3.cpp @@ -0,0 +1,55 @@ +#include "model/diffusion/minimax_h3.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_minimax_h3(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + true, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& minimax_h3_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MINIMAX_H3; + definition.geometry = {2, 24, 3}; + definition.name = "MiniMax-H3"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::MINIMAX_H3; + definition.traits.video_tae = true; + definition.sampling = {H3_AV_FLOW_PRED, 12.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_minimax_h3; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition = result.front(); + definition.version = VERSION_P_H3_IMAGE; + definition.name = "P-H3-Image"; + definition.traits.supports_video_generation = false; + definition.sampling.default_flow_shift = 1.f; + definition.ref_image_preset = "qwen"; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/minit2i.cpp b/src/model/integrations/minit2i.cpp new file mode 100644 index 0000000000..fb305943d9 --- /dev/null +++ b/src/model/integrations/minit2i.cpp @@ -0,0 +1,40 @@ +#include "model/diffusion/minit2i.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_minit2i(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model.model.net", + weight_manager); + return true; +} + +const std::vector& minit2i_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_MINIT2I; + definition.geometry = {1, 3, 3}; + definition.name = "MiniT2I"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {MINIT2I_FLOW_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_minit2i; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/moss_image.cpp b/src/model/integrations/moss_image.cpp new file mode 100644 index 0000000000..4a85d169a5 --- /dev/null +++ b/src/model/integrations/moss_image.cpp @@ -0,0 +1,33 @@ +#include "model_definition.h" + +#include "conditioning/conditioner.hpp" +#include "model/diffusion/moss_image.hpp" +#include "pipeline/model_builders.h" + +static bool build_moss_image(const sd::model_builders::Context& ctx, + const TokenizerConfig& tokenizers, + sd::model_builders::CoreRunners& runners) { + runners.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + ctx.tensor_storage_map, ctx.version, "", false, + ctx.weight_manager, tokenizers); + runners.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + ctx.tensor_storage_map, "model.diffusion_model", ctx.weight_manager); + return true; +} + +const std::vector& moss_image_definitions() { + static const std::vector definitions = [] { + ModelDefinition definition; + definition.version = VERSION_MOSS_IMAGE; + definition.geometry = {1, 128, 3}; + definition.name = "MOSS-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX2; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_moss_image; + return std::vector{std::move(definition)}; + }(); + return definitions; +} diff --git a/src/model/integrations/pid.cpp b/src/model/integrations/pid.cpp new file mode 100644 index 0000000000..d307fb6beb --- /dev/null +++ b/src/model/integrations/pid.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/pid.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_pid(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model.net", + weight_manager); + return true; +} + +const std::vector& pid_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_PID; + definition.geometry = {1, 3, 3}; + definition.name = "PiD"; + definition.traits.is_dit = true; + definition.traits.vae_format_override = true; + definition.sampling = {FLOW_PRED, 1.5f}; + definition.sampling.default_sample_method = LCM_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_pid; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/pixart.cpp b/src/model/integrations/pixart.cpp new file mode 100644 index 0000000000..d0a4aa7097 --- /dev/null +++ b/src/model/integrations/pixart.cpp @@ -0,0 +1,67 @@ +#include "model/diffusion/pixart.hpp" +#include "conditioning/conditioner.hpp" +#include "model/vae/auto_encoder_kl.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static void configure_pixart_vae(const sd::model_builders::Context& ctx, + const sd::model_builders::VAEOptions&, + VAE& vae) { + auto& model = static_cast(vae); + // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override. + if (ctx.tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) { + model.scale_factor = 0.18215f; + } + for (const auto& [key, value] : parse_key_value_args(ctx.params.model_args, "model arg")) { + if (key == "pixart_vae_scale_factor") { + float parsed = 0.f; + if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) { + model.scale_factor = parsed; + } else { + LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); + } + } + } + LOG_VERBOSE("pixart: VAE scale factor = %.5f", model.scale_factor); +} + +static bool build_pixart(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + false, + weight_manager, + sd_ctx_params->model_args); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); + return true; +} + +const std::vector& pixart_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_PIXART; + definition.geometry = {2, 4, 3}; + definition.name = "PixArt"; + definition.traits.is_dit = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_pixart; + definition.configure_vae = configure_pixart_vae; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/qwen_image.cpp b/src/model/integrations/qwen_image.cpp new file mode 100644 index 0000000000..49505bb779 --- /dev/null +++ b/src/model/integrations/qwen_image.cpp @@ -0,0 +1,87 @@ +#include "model/diffusion/qwen_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model/diffusion/qwen_image_2_1.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_qwen_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + enable_vision, + weight_manager, + tokenizers); + if (version == VERSION_QWEN_IMAGE_2_1) { + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager, + sd_ctx_params->model_args); + } else { + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager, + sd_ctx_params->model_args); + } + return true; +} + +const std::vector& qwen_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_QWEN_IMAGE; + definition.geometry = {1, 16, 3}; + definition.ref_image_preset = "qwen"; + definition.name = "Qwen Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_qwen_image; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_QWEN_IMAGE_LAYERED; + definition.geometry = {1, 16, 4}; + definition.ref_image_preset = "qwen_layered"; + definition.name = "Qwen Image Layered"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_qwen_image; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_QWEN_IMAGE_2_1; + definition.geometry = {2, 64, 4}; + definition.ref_image_preset = "qwen"; + definition.name = "Qwen Image 2.1"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = FLUX_SCHEDULER; + definition.build = build_qwen_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/sd.cpp b/src/model/integrations/sd.cpp new file mode 100644 index 0000000000..a36ff8e36d --- /dev/null +++ b/src/model/integrations/sd.cpp @@ -0,0 +1,261 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/unet.hpp" +#include "model/vae/auto_encoder_kl.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static void configure_sdxl_vae(const sd::model_builders::Context& ctx, + const sd::model_builders::VAEOptions& options, + VAE& vae) { + if (strlen(SAFE_STR(ctx.params.vae_path)) == 0 || ctx.params.force_sdxl_vae_conv_scale || options.external_vae_is_invalid) { + float vae_conv_2d_scale = 1.f / 32.f; + LOG_WARN( + "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, " + "using Conv2D scale %.3f", + vae_conv_2d_scale); + static_cast(vae).set_conv2d_scale(vae_conv_2d_scale); + } +} + +static bool build_sd(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + // SD1.x SD2.x SDXL + std::map embbeding_map; + for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) { + embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path)); + } + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + embbeding_map, + version, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} +static prediction_t detect_sdxl_prediction(const String2TensorStorage& weights) { + if (weights.count("edm_vpred.sigma_max")) + return EDM_V_PRED; + return weights.count("v_pred") ? V_PRED : EPS_PRED; +} + +const std::vector& sd_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_SD1; + definition.geometry = {8, 4, 3}; + definition.name = "SD 1.x"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD1_INPAINT; + definition.geometry = {8, 4, 3}; + definition.name = "SD 1.x Inpaint"; + definition.traits.is_inpaint = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD1_PIX2PIX; + definition.geometry = {8, 4, 3}; + definition.name = "Instruct-Pix2Pix"; + definition.traits.is_unet_edit = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD1_TINY_UNET; + definition.geometry = {8, 4, 3}; + definition.name = "SD 1.x Tiny UNet"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD2; + definition.geometry = {8, 4, 3}; + definition.name = "SD 2.x"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD2_INPAINT; + definition.geometry = {8, 4, 3}; + definition.name = "SD 2.x Inpaint"; + definition.traits.is_inpaint = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SD2_TINY_UNET; + definition.geometry = {8, 4, 3}; + definition.name = "SD 2.x Tiny UNet"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXS_512_DS; + definition.geometry = {8, 4, 3}; + definition.name = "SDXS (512-DS)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXS_09; + definition.geometry = {8, 4, 3}; + definition.name = "SDXS (09)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.probe_v_prediction = true; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_INPAINT; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL Inpaint"; + definition.traits.is_inpaint = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_PIX2PIX; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL Instruct-Pix2Pix"; + definition.traits.is_unet_edit = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_VEGA; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL (Vega)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SDXL_SSD1B; + definition.geometry = {8, 4, 3}; + definition.name = "SDXL (SSD1B)"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + definition.sampling.prediction_from_weights = detect_sdxl_prediction; + definition.configure_vae = configure_sdxl_vae; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_SVD; + definition.geometry = {8, 4, 3}; + definition.name = "SVD"; + definition.traits.supports_video_generation = true; + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_ESRGAN; + definition.geometry = {8, 4, 3}; + definition.name = "ESRGAN"; + + definition.sampling = {EPS_PRED, INFINITY}; + definition.sampling.default_sample_method = EULER_A_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/sd3.cpp b/src/model/integrations/sd3.cpp new file mode 100644 index 0000000000..f9c2ad7801 --- /dev/null +++ b/src/model/integrations/sd3.cpp @@ -0,0 +1,40 @@ +#include "conditioning/conditioner.hpp" +#include "model/diffusion/mmdit.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_sd3(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& sd3_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_SD3; + definition.geometry = {1, 16, 3}; + definition.name = "SD3.x"; + definition.traits.is_dit = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sd3; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/sensenova_u1.cpp b/src/model/integrations/sensenova_u1.cpp new file mode 100644 index 0000000000..900f222e1c --- /dev/null +++ b/src/model/integrations/sensenova_u1.cpp @@ -0,0 +1,38 @@ +#include "model/diffusion/sensenova_u1.h" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_sensenova_u1(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "", + weight_manager); + return true; +} + +const std::vector& sensenova_u1_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_SENSENOVA_U1_5; + definition.geometry = {32, 3, 3}; + definition.name = "SenseNova U1.5"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {SENSENOVA_U1_FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_sensenova_u1; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/wan.cpp b/src/model/integrations/wan.cpp new file mode 100644 index 0000000000..2daada2941 --- /dev/null +++ b/src/model/integrations/wan.cpp @@ -0,0 +1,119 @@ +#include "model/diffusion/wan.hpp" +#include "conditioning/conditioner.hpp" +#include "model/audio/wav2vec2.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_wan(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + true, + 0, + true, + weight_manager); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { + result.high_noise_diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.high_noise_diffusion_model", + version, + weight_manager); + } + if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" || + result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" || + result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") { + if (!sd::model_builders::ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) { + return false; + } + result.clip_vision = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION), + tensor_storage_map, + weight_manager); + } + if (version == VERSION_WAN2_2_S2V && + tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) { + if (!sd::model_builders::ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) { + return false; + } + result.audio_encoder = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER), + tensor_storage_map, + "wav2vec2.", + weight_manager); + } + return true; +} + +const std::vector& wan_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_WAN2; + definition.geometry = {2, 16, 3}; + definition.name = "Wan 2.x"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_WAN2_2_I2V; + definition.geometry = {2, 16, 3}; + definition.name = "Wan 2.2 I2V"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_WAN2_2_TI2V; + definition.geometry = {2, 48, 3}; + definition.name = "Wan 2.2 TI2V"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + { + ModelDefinition definition; + definition.version = VERSION_WAN2_2_S2V; + definition.geometry = {2, 16, 3}; + definition.name = "Wan 2.2 S2V"; + definition.traits.is_dit = true; + definition.traits.supports_video_generation = true; + definition.traits.vae_kind = ModelVAEKind::WAN; + definition.traits.video_tae = true; + definition.sampling = {FLOW_PRED, 5.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_wan; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/z_image.cpp b/src/model/integrations/z_image.cpp new file mode 100644 index 0000000000..29054d473a --- /dev/null +++ b/src/model/integrations/z_image.cpp @@ -0,0 +1,45 @@ +#include "model/diffusion/z_image.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_z_image(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + version, + weight_manager); + return true; +} + +const std::vector& z_image_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_Z_IMAGE; + definition.geometry = {1, 16, 3}; + definition.name = "Z-Image"; + definition.traits.is_dit = true; + definition.traits.vae_kind = ModelVAEKind::FLUX; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_z_image; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/integrations/z_image_l2p.cpp b/src/model/integrations/z_image_l2p.cpp new file mode 100644 index 0000000000..d8104f853a --- /dev/null +++ b/src/model/integrations/z_image_l2p.cpp @@ -0,0 +1,44 @@ +#include "model/diffusion/z_image_l2p.hpp" +#include "conditioning/conditioner.hpp" +#include "model_definition.h" +#include "pipeline/model_builders.h" + +static bool build_z_image_l2p(const sd::model_builders::Context& ctx, const TokenizerConfig& tokenizers, sd::model_builders::CoreRunners& result) { + const auto* sd_ctx_params = &ctx.params; + const auto& tensor_storage_map = ctx.tensor_storage_map; + const auto version = ctx.version; + const auto& weight_manager = ctx.weight_manager; + result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), + tensor_storage_map, + version, + "", + false, + weight_manager, + tokenizers); + result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), + tensor_storage_map, + "model.diffusion_model", + weight_manager); + return true; +} + +const std::vector& z_image_l2p_definitions() { + static const std::vector definitions = [] { + std::vector result; + { + ModelDefinition definition; + definition.version = VERSION_Z_IMAGE_L2P; + definition.geometry = {16, 3, 3}; + definition.name = "Z-Image L2P"; + definition.traits.is_dit = true; + definition.traits.fake_vae = true; + definition.sampling = {FLOW_PRED, 3.f}; + definition.sampling.default_sample_method = EULER_SAMPLE_METHOD; + definition.sampling.default_scheduler = DISCRETE_SCHEDULER; + definition.build = build_z_image_l2p; + result.push_back(std::move(definition)); + } + return result; + }(); + return definitions; +} diff --git a/src/model/te/llada_image_te.hpp b/src/model/te/llada_image_te.hpp index 1c97a4529a..2da78413b1 100644 --- a/src/model/te/llada_image_te.hpp +++ b/src/model/te/llada_image_te.hpp @@ -347,8 +347,8 @@ namespace LLaDAImageTE { blocks["visual.blocks." + std::to_string(i)] = std::make_shared(config); } blocks["vqmodel.quant_conv"] = std::make_shared(config.hidden_size, - config.codebook_embed_dim, - std::make_pair(1, 1)); + config.codebook_embed_dim, + std::make_pair(1, 1)); blocks["prior_projector.net.0.proj"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); blocks["prior_projector.net.2"] = std::make_shared(config.semantic_embed_dim, config.semantic_embed_dim, true); } diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 8ea918025e..add8b6ff9a 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -517,13 +517,13 @@ namespace LLM { auto up_proj = std::dynamic_pointer_cast(blocks["up_proj"]); auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); - auto h = gate_proj->forward(ctx, x); + auto h = gate_proj->forward(ctx, x); + auto up = up_proj->forward(ctx, x); if (activation == MLPActivation::GELU_TANH) { - h = ggml_ext_gelu(ctx->ggml_ctx, h, true); + h = ggml_geglu_split(ctx->ggml_ctx, h, up); } else { - h = ggml_silu_inplace(ctx->ggml_ctx, h); + h = ggml_swiglu_split(ctx->ggml_ctx, h, up); } - h = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x)); h = down_proj->forward(ctx, h); return h; } @@ -924,6 +924,8 @@ namespace LLM { int64_t txt_token_start = 0; int64_t txt_token_end = 0; ggml_tensor* input_embed = nullptr; + // Image and DeepStack indices remain in the expanded sequence coordinates. + int64_t removed_tokens = 0; for (int i = 0; i < image_embeds.size(); i++) { if (i == 0) { @@ -931,7 +933,10 @@ namespace LLM { } else { txt_token_start = image_embeds[i - 1].first + image_embeds[i - 1].second->ne[1]; } - txt_token_end = image_embeds[i].first; + txt_token_start -= removed_tokens; + txt_token_end = image_embeds[i].first - removed_tokens; + removed_tokens += image_embeds[i].second->ne[1] - 1; + GGML_ASSERT(txt_token_start >= 0 && txt_token_end >= txt_token_start && txt_token_end < raw_x->ne[1]); // An embed can sit flush against the previous one or at the very start/end of the // sequence, leaving no text tokens to splice around it. @@ -951,14 +956,14 @@ namespace LLM { } } - txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1]; + txt_token_start = image_embeds[image_embeds.size() - 1].first + image_embeds[image_embeds.size() - 1].second->ne[1] - removed_tokens; txt_token_end = raw_x->ne[1]; if (txt_token_end > txt_token_start) { auto final_txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); input_embed = ggml_concat(ctx->ggml_ctx, input_embed, final_txt_embed, 1); } - GGML_ASSERT(raw_x->ne[1] == input_embed->ne[1]); + GGML_ASSERT(raw_x->ne[1] + removed_tokens == input_embed->ne[1]); return input_embed; } @@ -1022,8 +1027,7 @@ namespace LLM { auto up_proj = std::dynamic_pointer_cast(blocks["up_proj"]); auto down_proj = std::dynamic_pointer_cast(blocks["down_proj"]); auto h = gate_proj->forward(ctx, x); - h = ggml_silu_inplace(ctx->ggml_ctx, h); - h = ggml_mul_inplace(ctx->ggml_ctx, h, up_proj->forward(ctx, x)); + h = ggml_swiglu_split(ctx->ggml_ctx, h, up_proj->forward(ctx, x)); x = down_proj->forward(ctx, h); } return x; @@ -1515,7 +1519,7 @@ namespace LLM { out = ggml_view_3d(c, out, head_dim, value_heads, n, head_dim * sizeof(float), head_dim * value_heads * sizeof(float), 0); out = std::dynamic_pointer_cast(blocks["norm"])->forward(ctx, out); auto z = ggml_reshape_3d(c, linear("in_proj_z"), head_dim, value_heads, n); - out = ggml_mul(c, out, ggml_silu(c, z)); + out = ggml_swiglu_split(c, z, out); out = ggml_reshape_3d(c, out, head_dim * value_heads, n, 1); return std::dynamic_pointer_cast(blocks["out_proj"])->forward(ctx, out); } @@ -2506,6 +2510,9 @@ namespace LLM { } int64_t n_tokens = input_ids->ne[0]; + for (const auto& image_embed : image_embeds) { + n_tokens += image_embed.second->ne[1] - 1; + } if (config.arch == LLMArch::MISTRAL_SMALL_3_2 || config.arch == LLMArch::MINISTRAL_3_3B || config.arch == LLMArch::QWEN3 || @@ -3014,20 +3021,17 @@ namespace LLM { LOG_VERBOSE("llm encode_image test done in %lldms", t1 - t0); } - std::string placeholder = "<|image_pad|>"; - std::string img_prompt = "Picture 1: <|vision_start|>"; // [24669, 220, 16, 25, 220, 151652] - int64_t num_image_tokens = image_embed.shape()[1]; - img_prompt.reserve(num_image_tokens * placeholder.size()); - for (int i = 0; i < num_image_tokens; i++) { - img_prompt += placeholder; - } - img_prompt += "<|vision_end|>"; - - std::vector>> image_embeds; - image_embeds.emplace_back(64, image_embed); + const std::string image_prefix = "Picture 1: <|vision_start|>"; + const std::string img_prompt = image_prefix + "<|image_pad|><|vision_end|>"; std::pair prompt_attn_range; std::string text = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n"; + std::vector prefix_tokens; + if (!tokenizer->encode(text + image_prefix, prefix_tokens, nullptr)) { + return; + } + std::vector>> image_embeds; + image_embeds.emplace_back(static_cast(prefix_tokens.size()), image_embed); text += img_prompt; prompt_attn_range.first = static_cast(text.size()); text += "change 'flux.cpp' to 'edit.cpp'"; diff --git a/src/model/te/ph3_image_connector.hpp b/src/model/te/ph3_image_connector.hpp new file mode 100644 index 0000000000..1e74e8e765 --- /dev/null +++ b/src/model/te/ph3_image_connector.hpp @@ -0,0 +1,70 @@ +#ifndef __SD_MODEL_TE_PH3_IMAGE_CONNECTOR_HPP__ +#define __SD_MODEL_TE_PH3_IMAGE_CONNECTOR_HPP__ + +#include "core/ggml_runner.h" +#include "model/common/ggml_block.hpp" + +struct PH3ImageTextConnector : public UnaryBlock { + static constexpr int64_t input_dim = 5 * 2560; + static constexpr int64_t output_dim = 5120; + + PH3ImageTextConnector() { + blocks["proj"] = std::make_shared(input_dim, output_dim, false, true, true); + } + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map = {}, + const std::string prefix = "") override { + for (const auto* name : {"input_mean", "input_std"}) { + params[name] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, input_dim); + } + for (const auto* name : {"output_mean", "output_std", "sink_embedding"}) { + params[name] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, output_dim); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto proj = std::dynamic_pointer_cast(blocks["proj"]); + x = ggml_div(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, x, params["input_mean"]), params["input_std"]); + x = proj->forward(ctx, x); + x = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, params["output_std"]), params["output_mean"]); + auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, 1)); + if (x->ne[1] == 1) { + return sink; + } + return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]), 1); + } +}; + +struct PH3ImageTextConnectorRunner : public GGMLRunner { + PH3ImageTextConnector model; + static constexpr const char* prefix = "text_encoders.ph3_image_connector"; + + PH3ImageTextConnectorRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager) { + model.init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "ph3_image_text_connector"; + } + + void get_param_tensors(std::map& tensors) { + model.get_param_tensors(tensors, prefix); + } + + sd::Tensor compute(int n_threads, const sd::Tensor& hidden_states) { + auto get_graph = [&]() { + auto gf = ggml_new_graph(compute_ctx); + auto x = make_input(hidden_states); + auto runner_ctx = get_context(); + ggml_build_forward_expand(gf, model.forward(&runner_ctx, x)); + return gf; + }; + return take_or_empty(GGMLRunner::compute(get_graph, n_threads, false)); + } +}; + +#endif // __SD_MODEL_TE_PH3_IMAGE_CONNECTOR_HPP__ diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index f445fe8005..7e7ac30984 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -156,9 +156,9 @@ struct T5DenseGatedActDense : public UnaryBlock { auto wi_1 = std::dynamic_pointer_cast(blocks["wi_1"]); auto wo = std::dynamic_pointer_cast(blocks["wo"]); - auto hidden_gelu = ggml_ext_gelu(ctx->ggml_ctx, wi_0->forward(ctx, x), true); + auto hidden_gate = wi_0->forward(ctx, x); auto hidden_linear = wi_1->forward(ctx, x); - x = ggml_mul_inplace(ctx->ggml_ctx, hidden_gelu, hidden_linear); + x = ggml_geglu_split(ctx->ggml_ctx, hidden_gate, hidden_linear); x = wo->forward(ctx, x); return x; } diff --git a/src/model/upscaler/h3_latent_upscaler.hpp b/src/model/upscaler/h3_latent_upscaler.hpp new file mode 100644 index 0000000000..006415081f --- /dev/null +++ b/src/model/upscaler/h3_latent_upscaler.hpp @@ -0,0 +1,584 @@ +#ifndef __SD_MODEL_UPSCALER_H3_LATENT_UPSCALER_HPP__ +#define __SD_MODEL_UPSCALER_H3_LATENT_UPSCALER_HPP__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_graph_cut.h" +#include "core/ggml_runner.h" +#include "core/tensor.hpp" +#include "core/util.h" +#include "model/common/ggml_block.hpp" +#include "model/upscaler/ltx_latent_upscaler.hpp" +#include "model_loader.h" + +// Ported from the community ComfyUI node LBH-123-AI/Comfyui_Minimax_h3_latent_Upscaler +// (MIT) with checkpoint LBH-123-AI/Minimax_h3_latent_Upscaler (Apache-2.0). + +namespace H3VUpsampler { + constexpr int H3_UPSAMPLER_GRAPH_SIZE = 8192; + + struct H3LatentUpscalerConfig { + int64_t in_channels = 24; + int64_t mid_channels = 512; + int embed_dim = 64; + int temporal_kernel = 5; + // Residual and temporal blocks share interleaved checkpoint indices. + std::vector in_res, in_temporal, out_res, out_temporal; + bool valid = false; + + static H3LatentUpscalerConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "") { + H3LatentUpscalerConfig config; + std::string base = prefix; + if (!base.empty()) { + base += "."; + } + auto find_weight = [&](const std::string& suffix) -> const TensorStorage* { + auto iter = tensor_storage_map.find(base + suffix); + if (iter == tensor_storage_map.end()) { + return nullptr; + } + return &iter->second; + }; + + const TensorStorage* conv_in = find_weight("conv_in.weight"); + if (conv_in == nullptr) { + LOG_ERROR("h3 latent upsampler: 'conv_in.weight' not found under prefix '%s'", prefix.c_str()); + return config; + } + // Bias shapes preserve channel counts when the loader folds 5D conv weights. + const TensorStorage* conv_in_bias = find_weight("conv_in.bias"); + if (conv_in_bias != nullptr) { + config.mid_channels = conv_in_bias->ne[0]; + } else if (conv_in->n_dims == 5) { + config.mid_channels = conv_in->ne[0]; + } + const TensorStorage* conv_out_bias = find_weight("conv_out.bias"); + if (conv_out_bias != nullptr) { + config.in_channels = conv_out_bias->ne[0]; + } else if (conv_in->n_dims == 5) { + config.in_channels = conv_in->ne[1]; + } + if (config.mid_channels <= 0 || config.in_channels <= 0) { + LOG_ERROR("h3 latent upsampler: could not infer channel counts from the checkpoint"); + return config; + } + + const TensorStorage* embed_mid = find_weight("embed.2.weight"); + if (embed_mid != nullptr && embed_mid->n_dims >= 1) { + config.embed_dim = static_cast(embed_mid->ne[0]); + } + + std::set in_res_idx, in_temporal_idx, out_res_idx, out_temporal_idx; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, base)) { + continue; + } + const std::string rel = name.substr(base.size()); + const bool is_in = starts_with(rel, "in_blocks."); + const bool is_out = starts_with(rel, "out_blocks."); + if (!is_in && !is_out) { + continue; + } + const std::string rest = rel.substr(is_in ? 10 : 11); + const size_t dot = rest.find('.'); + if (dot == std::string::npos) { + continue; + } + int idx = 0; + try { + idx = std::stoi(rest.substr(0, dot)); + } catch (...) { + continue; + } + const std::string part = rest.substr(dot + 1); + if (starts_with(part, "dwconv.") || starts_with(part, "pwconv.")) { + (is_in ? in_temporal_idx : out_temporal_idx).insert(idx); + } else if (starts_with(part, "in_layers.") || starts_with(part, "out_norm.") || + starts_with(part, "out_layers.") || starts_with(part, "emb_layers.")) { + (is_in ? in_res_idx : out_res_idx).insert(idx); + } else if (starts_with(part, "attn.")) { + LOG_ERROR("h3 latent upsampler: attention blocks are not supported by the v1 runner"); + return config; + } + } + config.in_res = std::vector(in_res_idx.begin(), in_res_idx.end()); + config.in_temporal = std::vector(in_temporal_idx.begin(), in_temporal_idx.end()); + config.out_res = std::vector(out_res_idx.begin(), out_res_idx.end()); + config.out_temporal = std::vector(out_temporal_idx.begin(), out_temporal_idx.end()); + if (config.in_res.empty() || config.out_res.empty()) { + LOG_ERROR("h3 latent upsampler: no res blocks found in the checkpoint"); + return config; + } + + auto temporal_kernel_of = [&](std::vector idxs, const char* stage) -> int { + for (int idx : idxs) { + const TensorStorage* dw = find_weight(std::string(stage) + "." + std::to_string(idx) + ".dwconv.weight"); + if (dw != nullptr) { + // kD stays at ne[2] in both [OC, IC, kD, kH, kW] and [kW, kH, kD, OC*IC]. + return static_cast(dw->ne[2]); + } + } + return 0; + }; + const int k_in = temporal_kernel_of(config.in_temporal, "in_blocks"); + const int k_out = k_in > 0 ? k_in : temporal_kernel_of(config.out_temporal, "out_blocks"); + if (k_out > 0) { + config.temporal_kernel = k_out; + } else { + config.temporal_kernel = 1; + } + + config.valid = true; + LOG_VERBOSE("h3 latent upsampler: in_channels = %" PRId64 ", mid_channels = %" PRId64 + ", in_res = %zu, in_temporal = %zu, out_res = %zu, out_temporal = %zu, temporal_kernel = %d, embed_dim = %d", + config.in_channels, + config.mid_channels, + config.in_res.size(), + config.in_temporal.size(), + config.out_res.size(), + config.out_temporal.size(), + config.temporal_kernel, + config.embed_dim); + return config; + } + }; + + class H3ResBlock : public GGMLBlock { + protected: + int64_t channels; + int64_t embed_dim; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map, + const std::string prefix) override { + auto make_conv = [&](const std::string& name) { + blocks[name] = std::shared_ptr(new Conv3d(channels, + channels, + {3, 3, 3}, + {1, 1, 1}, + {1, 1, 1})); + }; + make_conv("in_layers.2"); + make_conv("out_layers.2"); + blocks["in_layers.0"] = std::shared_ptr(new LTXVUpsampler::VideoGroupNorm(32, channels)); + blocks["out_norm"] = std::shared_ptr(new LTXVUpsampler::VideoGroupNorm(32, channels)); + blocks["emb_layers.1"] = std::shared_ptr(new Linear(embed_dim, 2 * channels, true)); + } + + public: + H3ResBlock(int64_t channels, int64_t embed_dim) + : channels(channels), + embed_dim(embed_dim) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb) { + auto in_layers = std::dynamic_pointer_cast(blocks["in_layers.2"]); + auto in_norm = std::dynamic_pointer_cast(blocks["in_layers.0"]); + auto out_layers = std::dynamic_pointer_cast(blocks["out_layers.2"]); + auto out_norm = std::dynamic_pointer_cast(blocks["out_norm"]); + auto emb_layers = std::dynamic_pointer_cast(blocks["emb_layers.1"]); + + ggml_tensor* h = in_norm->forward(ctx, x); + h = ggml_silu_inplace(ctx->ggml_ctx, h); + h = in_layers->forward(ctx, h); + + // Keep emb intact: every residual block uses the same embedding. + ggml_tensor* emb_silu = ggml_silu(ctx->ggml_ctx, emb); + ggml_tensor* emb_out = emb_layers->forward(ctx, emb_silu); + ggml_tensor* emb_1d = ggml_reshape_1d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, emb_out), 2 * channels); + ggml_tensor* scale = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels); + ggml_tensor* shift = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, emb_1d, 0, channels, 2 * channels), 1, 1, 1, channels); + ggml_tensor* gain = ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f); + + h = out_norm->forward(ctx, h); + h = ggml_mul_inplace(ctx->ggml_ctx, h, gain); + h = ggml_add_inplace(ctx->ggml_ctx, h, shift); + h = ggml_silu_inplace(ctx->ggml_ctx, h); + h = out_layers->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, h, x); + } + }; + + // Shifted multiply-adds preserve depthwise weights [1, 1, k, C] without + // expanding them into the dense kernels required by Conv3d. + class H3TemporalDWConv3d : public UnaryBlock { + protected: + int64_t channels; + int kernel; + bool has_bias; + std::string prefix; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map, + const std::string prefix) override { + this->prefix = prefix; + params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, 1, kernel, channels); + if (has_bias) { + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + } + } + + public: + H3TemporalDWConv3d(int64_t channels, int kernel, bool has_bias = true) + : channels(channels), + kernel(kernel), + has_bias(has_bias) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + ggml_tensor* w = params["weight"]; + if (ctx->weight_adapter) { + w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + } + ggml_tensor* w32 = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32); + ggml_tensor* xa = x; + if (xa->type != GGML_TYPE_F32) { + xa = ggml_cast(ctx->ggml_ctx, xa, GGML_TYPE_F32); + } + + const int64_t T = xa->ne[2]; + const int pad = kernel / 2; + + // ggml_pad only pads the end; concatenate zeros for symmetric padding. + ggml_tensor* xp = xa; + if (pad > 0) { + ggml_tensor* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, xa, 2, 0, 1), 0.0f); + ggml_tensor* lead = zero; + for (int i = 1; i < pad; ++i) { + lead = ggml_concat(ctx->ggml_ctx, lead, zero, 2); + } + xp = ggml_concat(ctx->ggml_ctx, lead, xa, 2); + xp = ggml_concat(ctx->ggml_ctx, xp, lead, 2); + } + + ggml_tensor* acc = nullptr; + for (int j = 0; j < kernel; ++j) { + ggml_tensor* xs = ggml_ext_slice(ctx->ggml_ctx, xp, 2, j, j + T); + ggml_tensor* wj = ggml_ext_slice(ctx->ggml_ctx, w32, 2, j, j + 1); // [1, 1, 1, C] + ggml_tensor* term = ggml_mul(ctx->ggml_ctx, xs, wj); + acc = acc == nullptr ? term : ggml_add(ctx->ggml_ctx, acc, term); + } + if (has_bias) { + ggml_tensor* b = params["bias"]; + if (ctx->weight_adapter) { + b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + } + acc = ggml_add(ctx->ggml_ctx, acc, ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, 1, channels)); + } + return acc; + } + }; + + class H3TemporalConv : public GGMLBlock { + protected: + int64_t channels; + int temporal_kernel; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map, + const std::string prefix) override { + blocks["dwconv"] = std::shared_ptr(new H3TemporalDWConv3d(channels, temporal_kernel)); + blocks["pwconv"] = std::shared_ptr(new Conv3d(channels, + channels, + {1, 1, 1}, + {1, 1, 1}, + {0, 0, 0})); + blocks["norm"] = std::shared_ptr(new LTXVUpsampler::VideoGroupNorm(32, channels)); + } + + public: + H3TemporalConv(int64_t channels, int temporal_kernel) + : channels(channels), + temporal_kernel(temporal_kernel) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto dwconv = std::dynamic_pointer_cast(blocks["dwconv"]); + auto pwconv = std::dynamic_pointer_cast(blocks["pwconv"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + + ggml_tensor* h = norm->forward(ctx, x); + h = ggml_silu_inplace(ctx->ggml_ctx, h); + h = dwconv->forward(ctx, h); + h = pwconv->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, h, x); + } + }; + + class H3Embed : public GGMLBlock { + protected: + int embed_dim; + + void init_params(ggml_context* ctx, + const String2TensorStorage& tensor_storage_map, + const std::string prefix) override { + blocks["0"] = std::shared_ptr(new Linear(1, embed_dim, true)); + blocks["2"] = std::shared_ptr(new Linear(embed_dim, embed_dim, true)); + } + + public: + explicit H3Embed(int embed_dim) + : embed_dim(embed_dim) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* scale_input) { + auto l0 = std::dynamic_pointer_cast(blocks["0"]); + auto l2 = std::dynamic_pointer_cast(blocks["2"]); + ggml_tensor* h = l0->forward(ctx, ggml_reshape_2d(ctx->ggml_ctx, scale_input, 1, 1)); + h = ggml_silu_inplace(ctx->ggml_ctx, h); + return l2->forward(ctx, h); + } + }; + + class H3LatentUpscaler : public GGMLBlock { + public: + H3LatentUpscalerConfig config; + + explicit H3LatentUpscaler(H3LatentUpscalerConfig config_) + : config(std::move(config_)) { + for (int idx : config.in_res) { + blocks["in_blocks." + std::to_string(idx)] = std::shared_ptr(new H3ResBlock(config.mid_channels, + config.embed_dim)); + } + for (int idx : config.in_temporal) { + blocks["in_blocks." + std::to_string(idx)] = std::shared_ptr(new H3TemporalConv(config.mid_channels, + config.temporal_kernel)); + } + for (int idx : config.out_res) { + blocks["out_blocks." + std::to_string(idx)] = std::shared_ptr(new H3ResBlock(config.mid_channels, + config.embed_dim)); + } + for (int idx : config.out_temporal) { + blocks["out_blocks." + std::to_string(idx)] = std::shared_ptr(new H3TemporalConv(config.mid_channels, + config.temporal_kernel)); + } + blocks["conv_in"] = std::shared_ptr(new Conv3d(config.in_channels, + config.mid_channels, + {3, 3, 3}, + {1, 1, 1}, + {1, 1, 1})); + blocks["norm_out"] = std::shared_ptr(new LTXVUpsampler::VideoGroupNorm(32, config.mid_channels)); + blocks["conv_out"] = std::shared_ptr(new Conv3d(config.mid_channels, + config.in_channels, + {3, 3, 3}, + {1, 1, 1}, + {1, 1, 1})); + blocks["embed"] = std::shared_ptr(new H3Embed(config.embed_dim)); + } + + ggml_tensor* forward_stage_a(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* scale_input) { + auto conv_in = std::dynamic_pointer_cast(blocks["conv_in"]); + auto embed = std::dynamic_pointer_cast(blocks["embed"]); + + ggml_tensor* h = conv_in->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(h, "h3_latent_upscaler.conv_in", "x"); + + ggml_tensor* emb = embed->forward(ctx, scale_input); + std::set temporal_set(config.in_temporal.begin(), config.in_temporal.end()); + int max_idx = 0; + for (int idx : config.in_res) { + max_idx = std::max(max_idx, idx); + } + for (int idx : config.in_temporal) { + max_idx = std::max(max_idx, idx); + } + for (int i = 0; i <= max_idx; ++i) { + if (temporal_set.count(i)) { + auto block = std::dynamic_pointer_cast(blocks["in_blocks." + std::to_string(i)]); + h = block->forward(ctx, h); + } else { + auto block = std::dynamic_pointer_cast(blocks["in_blocks." + std::to_string(i)]); + h = block->forward(ctx, h, emb); + } + sd::ggml_graph_cut::mark_graph_cut(h, "h3_latent_upscaler.in_blocks." + std::to_string(i), "x"); + } + return h; + } + + ggml_tensor* forward_stage_b(GGMLRunnerContext* ctx, ggml_tensor* m, ggml_tensor* scale_input) { + auto conv_out = std::dynamic_pointer_cast(blocks["conv_out"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto embed = std::dynamic_pointer_cast(blocks["embed"]); + + ggml_tensor* g = m; + ggml_tensor* emb = embed->forward(ctx, scale_input); + std::set temporal_set(config.out_temporal.begin(), config.out_temporal.end()); + int max_idx = 0; + for (int idx : config.out_res) { + max_idx = std::max(max_idx, idx); + } + for (int idx : config.out_temporal) { + max_idx = std::max(max_idx, idx); + } + for (int i = 0; i <= max_idx; ++i) { + if (temporal_set.count(i)) { + auto block = std::dynamic_pointer_cast(blocks["out_blocks." + std::to_string(i)]); + g = block->forward(ctx, g); + } else { + auto block = std::dynamic_pointer_cast(blocks["out_blocks." + std::to_string(i)]); + g = block->forward(ctx, g, emb); + } + sd::ggml_graph_cut::mark_graph_cut(g, "h3_latent_upscaler.out_blocks." + std::to_string(i), "x"); + } + + g = norm_out->forward(ctx, g); + g = ggml_silu_inplace(ctx->ggml_ctx, g); + g = conv_out->forward(ctx, g); + sd::ggml_graph_cut::mark_graph_cut(g, "h3_latent_upscaler.conv_out", "x"); + return g; + } + }; + + struct H3LatentUpscalerRunner : public GGMLRunner { + H3LatentUpscalerConfig config; + std::unique_ptr model; + std::string weight_prefix; + + H3LatentUpscalerRunner(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "", + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(H3LatentUpscalerConfig::detect_from_weights(tensor_storage_map, prefix)), + weight_prefix(prefix) { + if (!config.valid || config.embed_dim <= 0 || config.mid_channels <= 0) { + LOG_ERROR("unsupported H3 latent upsampler config: in_channels=%" PRId64 " mid_channels=%" PRId64 " embed_dim=%d temporal_kernel=%d", + config.in_channels, + config.mid_channels, + config.embed_dim, + config.temporal_kernel); + return; + } + model = std::make_unique(config); + model->init(params_ctx, tensor_storage_map, prefix); + } + + std::string get_desc() override { + return "h3_latent_upscaler"; + } + + void get_param_tensors(std::map& tensors) { + if (model) { + model->get_param_tensors(tensors, weight_prefix); + } + } + + // x: H3 diffusion latent [W, H, T, C] (or 5D batch 1). The output stays + // in the same normalized domain for the next sampling pass. + sd::Tensor compute(const int n_threads, + const sd::Tensor& x, + int64_t target_w, + int64_t target_h) { + if (model == nullptr) { + LOG_ERROR("H3 latent upsampler is not loaded"); + return {}; + } + if (x.dim() != 4 && x.dim() != 5) { + LOG_ERROR("H3 latent upsampler expects a 4D or 5D video latent, got dim=%lld", (long long)x.dim()); + return {}; + } + if (x.dim() == 5 && x.shape()[4] != 1) { + LOG_ERROR("H3 latent upsampler currently supports batch size 1, got batch=%lld", (long long)x.shape()[4]); + return {}; + } + if (config.in_channels != 24 || x.shape()[3] != config.in_channels) { + LOG_ERROR("H3 latent upsampler requires 24 channels, got model=%" PRId64 " input=%lld", config.in_channels, (long long)x.shape()[3]); + return {}; + } + const int64_t in_w = x.shape()[0]; + const int64_t in_h = x.shape()[1]; + const int64_t T = x.shape()[2]; + if (target_w % 2 != 0 || target_h % 2 != 0) { + LOG_ERROR("H3 latent upsampler target latent size must be even (pixel size a multiple of 32), got %" PRId64 "x%" PRId64, + target_w, + target_h); + return {}; + } + const float scale = (static_cast(target_w) / static_cast(in_w) + + static_cast(target_h) / static_cast(in_h)) / + 2.0f; + if (scale < 1.0f - 1e-3f) { + LOG_ERROR("H3 latent upsampler only supports upscaling (scale >= 1.0), got %.3f", scale); + return {}; + } + + const size_t expected_dim = static_cast(x.dim()); + const sd::Tensor mean({1, 1, 1, 24}, + {0.858090341091156f, -0.9606591463088989f, 1.0661640167236328f, -0.5090325474739075f, + -0.2727581858634949f, -1.3675414323806763f, -0.2553254961967468f, -0.26907554268836975f, + -0.5376840829849243f, -0.0464097298681736f, 0.6657370328903198f, 0.19690127670764923f, + -0.5460608005523682f, -0.4035342037677765f, -0.23683024942874908f, 0.25928452610969543f, + -0.30133944749832153f, 0.211341992020607f, -1.1206848621368408f, 0.3581933379173279f, + -0.04225143790245056f, 0.2604829967021942f, 0.22864092886447906f, 0.7056031823158264f}); + const sd::Tensor stddev({1, 1, 1, 24}, + {1.2223774194717407f, 1.2767263650894165f, 1.6831774711608887f, 1.7549455165863037f, + 1.5636216402053833f, 2.194143533706665f, 0.9653137922286987f, 1.0569885969161987f, + 0.841948926448822f, 0.7729952931404114f, 1.8955937623977661f, 0.946841835975647f, + 0.7996809482574463f, 0.44988900423049927f, 0.7197399735450745f, 0.6936293244361877f, + 2.961095094680786f, 2.7694199085235596f, 3.0496184825897217f, 2.1088054180145264f, + 3.276226282119751f, 3.1627357006073f, 2.2816812992095947f, 2.6127843856811523f}); + // Training applies the VAE normalization a second time to diffusion latents. + // Undo only this extra transform after the network. + const sd::Tensor normalized = (x - mean) / stddev; + const sd::Tensor scale_tensor = sd::Tensor::from_vector({scale - 1.0f}); + auto get_graph_a = [this, &normalized, &scale_tensor]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(H3_UPSAMPLER_GRAPH_SIZE); + ggml_tensor* x_t = make_input(normalized); + ggml_tensor* scale_input = make_input(scale_tensor); + auto runner_ctx = get_context(gf); + ggml_tensor* out = model->forward_stage_a(&runner_ctx, x_t, scale_input); + ggml_build_forward_expand(gf, out); + return gf; + }; + std::optional> mid_opt = GGMLRunner::compute(get_graph_a, n_threads, false, false); + runner_end(); + if (!mid_opt.has_value() || mid_opt->empty()) { + LOG_ERROR("H3 latent upsampler stage A (in-block chain) failed"); + return {}; + } + sd::Tensor mid = std::move(*mid_opt); + + std::vector mid_shape = mid.shape(); + mid_shape[0] = target_w; + mid_shape[1] = target_h; + sd::Tensor mid_up = sd::ops::interpolate(mid, mid_shape, sd::ops::InterpolateMode::Bilinear, false, false); + if (mid_up.empty()) { + LOG_ERROR("H3 latent upsampler middle spatial interpolate failed: %" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64, + in_w, in_h, target_w, target_h); + return {}; + } + + auto get_graph_b = [this, &mid_up, &scale_tensor]() -> ggml_cgraph* { + ggml_cgraph* gf = new_graph_custom(H3_UPSAMPLER_GRAPH_SIZE); + ggml_tensor* m_t = make_input(mid_up); + ggml_tensor* scale_input = make_input(scale_tensor); + auto runner_ctx = get_context(gf); + ggml_tensor* out = model->forward_stage_b(&runner_ctx, m_t, scale_input); + ggml_build_forward_expand(gf, out); + return gf; + }; + std::optional> out_opt = GGMLRunner::compute(get_graph_b, n_threads, false, false); + runner_end(); + if (!out_opt.has_value() || out_opt->empty()) { + LOG_ERROR("H3 latent upsampler stage B (out-block chain) failed"); + return {}; + } + sd::Tensor out = std::move(*out_opt); + if (out.shape()[0] != target_w || out.shape()[1] != target_h || out.shape()[2] != T) { + LOG_ERROR("H3 latent upsampler output shape mismatch: got %" PRId64 "x%" PRId64 "x%" PRId64 " expected %" PRId64 "x%" PRId64 "x%" PRId64, + out.shape()[0], out.shape()[1], out.shape()[2], target_w, target_h, T); + return {}; + } + return restore_trailing_singleton_dims(out * stddev + mean, expected_dim); + } + }; + +} // namespace H3VUpsampler + +#endif // __SD_MODEL_UPSCALER_H3_LATENT_UPSCALER_HPP__ diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index 2ff8528ecc..020392149b 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -1170,9 +1170,9 @@ namespace LTXV { auto& tensor_storage_map = model_loader.get_tensor_storage_map(); auto ltx_audio_vae = std::make_shared(backend, - tensor_storage_map, - prefix, - model_manager); + tensor_storage_map, + prefix, + model_manager); if (!model_manager->set_loader(std::move(model_loader)) || !model_manager->register_runner_params(ModelComponent::AudioVAE, diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index f680f20081..3fa42a1662 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -497,12 +497,12 @@ namespace LTXVAE { GGML_ASSERT(out_channels % factor == 0); blocks["conv"] = std::make_shared(in_channels, - out_channels / factor, - 3, - std::tuple{1, 1, 1}, - 1, - true, - force_conv_prec_f32); + out_channels / factor, + 3, + std::tuple{1, 1, 1}, + 1, + true, + force_conv_prec_f32); blocks["skip_downsample"] = std::make_shared(in_channels, out_channels, factor_t, factor_s); blocks["conv_downsample"] = std::make_shared(out_channels / factor, out_channels, factor_t, factor_s); } @@ -1077,11 +1077,11 @@ namespace LTXVAE { patch_size); } blocks["decoder"] = std::make_shared(version, - tensor_storage_map, - prefix, - patch_size, - false, - timestep_conditioning); + tensor_storage_map, + prefix, + patch_size, + false, + timestep_conditioning); blocks["per_channel_statistics"] = std::make_shared(); } diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index c150559477..c3ba2036a1 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -86,10 +86,10 @@ namespace MiniMaxH3 { blocks["block.2"] = std::make_shared(in_channels, 9); blocks["block.3"] = std::make_shared(in_channels); blocks["block.4"] = std::make_shared(in_channels, - out_channels, - 2 * stride, - stride, - static_cast(std::ceil(stride / 2.f))); + out_channels, + 2 * stride, + stride, + static_cast(std::ceil(stride / 2.f))); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { @@ -144,8 +144,8 @@ namespace MiniMaxH3 { auto w1 = std::dynamic_pointer_cast(blocks["w1"]); auto w2 = std::dynamic_pointer_cast(blocks["w2"]); x = norm->forward(ctx, x); - auto gate = ggml_ext_gelu(ctx->ggml_ctx, w0->forward(ctx, x), true); - return w2->forward(ctx, ggml_mul(ctx->ggml_ctx, gate, w1->forward(ctx, x))); + auto gate = w0->forward(ctx, x); + return w2->forward(ctx, ggml_geglu_split(ctx->ggml_ctx, gate, w1->forward(ctx, x))); } }; @@ -310,12 +310,12 @@ namespace MiniMaxH3 { } blocks["activation_post"] = std::make_shared(channels, audio_direct_dw()); blocks["conv_post"] = std::make_shared(channels, - 1, - 7, - 1, - 3, - 1, - false); + 1, + 7, + 1, + 3, + 1, + false); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 2b1085d6be..d6a6bb06b7 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -14,6 +14,7 @@ #include #include +#include "core/rng_mt19937.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/vae/vae.hpp" @@ -23,6 +24,34 @@ namespace MiniMaxH3VAE { constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144; + struct MiniMaxH3VAEConfig { + int num_layers = 36; + int patch_size_t = 4; + + static MiniMaxH3VAEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { + MiniMaxH3VAEConfig config; + auto proj = tensors.find(prefix + ".decoder.proj_out.weight"); + if (proj != tensors.end()) { + config.patch_size_t = static_cast(proj->second.ne[1] / (3 * 16 * 16)); + if (config.patch_size_t != 1 && config.patch_size_t != 4) { + throw std::runtime_error("unsupported MiniMax-H3 VAE temporal patch size"); + } + int layers = 0; + const auto layer_prefix = prefix + ".decoder.transformer_blocks."; + for (const auto& [name, _] : tensors) { + if (starts_with(name, layer_prefix)) { + layers = std::max(layers, std::atoi(name.c_str() + layer_prefix.size()) + 1); + } + } + if (layers > 0) { + config.num_layers = layers; + } + LOG_VERBOSE("minimax_h3_vae: layers=%d, patch_size_t=%d", config.num_layers, config.patch_size_t); + } + return config; + } + }; + struct CausalConv3d : public Conv3d { std::tuple temporal_padding; @@ -453,18 +482,13 @@ namespace MiniMaxH3VAE { blocks["w2"] = std::make_shared(kInnerDim, dim, true); } + // fused_glu is kept for the fast-path callers; upstream (#2250) now always uses the native GLU. ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, bool fused_glu = false) { - auto w1 = std::dynamic_pointer_cast(blocks["w1"]); - auto w2 = std::dynamic_pointer_cast(blocks["w2"]); - if (fused_glu) { - // silu(first half) * second half in one op; same per-element math as silu + mul. - return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, w1->forward(ctx, x))); - } - auto gate = ggml_ext_chunk(ctx->ggml_ctx, w1->forward(ctx, x), 2, 0); - return w2->forward(ctx, - ggml_mul(ctx->ggml_ctx, - ggml_silu(ctx->ggml_ctx, gate[0]), - gate[1])); + SD_UNUSED(fused_glu); + auto w1 = std::dynamic_pointer_cast(blocks["w1"]); + auto w2 = std::dynamic_pointer_cast(blocks["w2"]); + auto gate = ggml_swiglu(ctx->ggml_ctx, w1->forward(ctx, x)); + return w2->forward(ctx, gate); } }; @@ -511,13 +535,14 @@ namespace MiniMaxH3VAE { }; struct Decoder : public GGMLBlock { - static constexpr int dim = 2048; - static constexpr int num_layers = 36; + static constexpr int dim = 2048; + int num_layers; static constexpr int num_register_tokens = 4; static constexpr int patch_size = 16; - static constexpr int patch_size_t = 4; + int patch_size_t; - Decoder() { + explicit Decoder(const MiniMaxH3VAEConfig& config) + : num_layers(config.num_layers), patch_size_t(config.patch_size_t) { blocks["x_embedder"] = std::make_shared(24, dim, true); for (int i = 0; i < num_layers; ++i) { blocks["transformer_blocks." + std::to_string(i)] = @@ -539,7 +564,9 @@ namespace MiniMaxH3VAE { GGML_TYPE_F32, dim, num_register_tokens); - params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim); + if (patch_size_t != 1) { + params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim); + } } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -624,7 +651,10 @@ namespace MiniMaxH3VAE { }; struct MiniMaxH3VideoVAE : public GGMLBlock { - MiniMaxH3VideoVAE() { + bool single_frame; + + explicit MiniMaxH3VideoVAE(const MiniMaxH3VAEConfig& config) + : single_frame(config.patch_size_t == 1) { blocks["encoder"] = std::make_shared(); blocks["quant_conv"] = std::make_shared(48, 48, @@ -632,7 +662,7 @@ namespace MiniMaxH3VAE { blocks["post_quant_conv"] = std::make_shared(24, 24, std::tuple{1, 1, 1}); - blocks["decoder"] = std::make_shared(); + blocks["decoder"] = std::make_shared(config); } ggml_tensor* encode(GGMLRunnerContext* ctx, @@ -645,7 +675,7 @@ namespace MiniMaxH3VAE { auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); auto quant = std::dynamic_pointer_cast(blocks["quant_conv"]); auto moments = quant->forward(ctx, encoder->forward(ctx, pixels)); - return ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24); + return single_frame ? moments : ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24); } ggml_tensor* decode(GGMLRunnerContext* ctx, @@ -666,6 +696,7 @@ namespace MiniMaxH3VAE { }; struct MiniMaxH3VideoVAERunner : public VAE { + MiniMaxH3VAEConfig config; MiniMaxH3VideoVAE model; sd::Tensor pixel_mean; sd::Tensor pixel_std; @@ -681,6 +712,8 @@ namespace MiniMaxH3VAE { const std::string& prefix = "first_stage_model", std::shared_ptr weight_manager = nullptr) : VAE(VERSION_MINIMAX_H3, backend, prefix, weight_manager), + config(MiniMaxH3VAEConfig::detect_from_weights(tensor_storage_map, prefix)), + model(config), pixel_mean({1, 1, 1, 3}, {0.485f, 0.456f, 0.406f}), pixel_std({1, 1, 1, 3}, {0.229f, 0.224f, 0.225f}), latents_mean({1, 1, 1, 24}, @@ -702,12 +735,12 @@ namespace MiniMaxH3VAE { } std::string get_desc() override { - return "minimax_h3_video_vae"; + return config.patch_size_t == 1 ? "p_h3_image_vae" : "minimax_h3_video_vae"; } int get_encoder_output_channels(int input_channels) override { SD_UNUSED(input_channels); - return 24; + return config.patch_size_t == 1 ? 48 : 24; } void get_param_tensors(std::map& tensors) override { @@ -717,6 +750,15 @@ namespace MiniMaxH3VAE { sd::Tensor vae_output_to_latents(const sd::Tensor& vae_output, std::shared_ptr rng) override { SD_UNUSED(rng); + if (config.patch_size_t == 1) { + auto mean = sd::ops::slice(vae_output, 3, 0, 24); + auto logvar = sd::ops::slice(vae_output, 3, 24, 48); + auto noise = MT19937RNG(42).randn(static_cast(mean.numel())); + for (int64_t i = 0; i < mean.numel(); ++i) { + mean[i] += std::exp(0.5f * std::clamp(logvar[i], -30.f, 20.f)) * noise[i]; + } + return mean; + } return vae_output; } @@ -809,6 +851,10 @@ namespace MiniMaxH3VAE { bool circular_y = false) override { auto input = ensure_video_shape(x); auto tiling = resolve_tiling_params(tiling_params); + if (config.patch_size_t == 1 && input.shape()[2] != 1) { + LOG_ERROR("P-H3-Image VAE encodes single-frame images only"); + return {}; + } if (input.shape()[2] == 1) { auto encoded = VAE::encode(n_threads, input, tiling, circular_x, circular_y); if (!encoded.empty() && encoded.shape()[2] > 1) { @@ -853,6 +899,10 @@ namespace MiniMaxH3VAE { bool silent = false) override { auto input = ensure_video_shape(x); auto tiling = resolve_tiling_params(tiling_params); + if (config.patch_size_t == 1 && input.shape()[2] != 1) { + LOG_ERROR("P-H3-Image VAE decodes single-frame latents only"); + return {}; + } if (input.shape()[2] == 1) { auto decoded = VAE::decode(n_threads, input, @@ -1224,7 +1274,9 @@ namespace MiniMaxH3VAE { ids.reserve(static_cast(width * height * num_frames + 5)); constexpr float two_pi = 6.28318530717958647692f; for (int64_t t = 0; t < num_frames; ++t) { - float pt = (2.f * ((t + 0.5f) / num_frames) - 1.f) * two_pi; + float pt = (config.patch_size_t == 1 ? 2.f * (0.5f / 7.f) - 1.f + : 2.f * ((t + 0.5f) / num_frames) - 1.f) * + two_pi; for (int64_t h = 0; h < height; ++h) { float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi; for (int64_t w = 0; w < width; ++w) { diff --git a/src/model/vae/vae_tiling.hpp b/src/model/vae/vae_tiling.hpp index 50cf08355e..acad046236 100644 --- a/src/model/vae/vae_tiling.hpp +++ b/src/model/vae/vae_tiling.hpp @@ -15,8 +15,9 @@ enum class VAETemporalDirection { }; struct VAETemporalTilingConfig { - int tile_frames = 1; - int overlap = 0; + int tile_frames = 1; + int overlap = 0; + int decode_frames = 1; }; struct VAETemporalTile { @@ -43,7 +44,8 @@ inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tilin config.overlap = std::max(0, default_overlap); for (const auto& [key, value] : parse_key_value_args(params.extra_tiling_args, "VAE extra tiling arg")) { - if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap") { + if (key != "temporal_tile_frames" && key != "temporal_tile_size" && key != "temporal_tile_overlap" && + key != "temporal_decode_frames") { continue; } @@ -52,6 +54,8 @@ inline VAETemporalTilingConfig resolve_vae_temporal_tiling_config(const sd_tilin LOG_WARN("ignoring invalid VAE extra tiling arg '%s=%s'", key.c_str(), value.c_str()); } else if (key == "temporal_tile_overlap") { config.overlap = std::max(0, parsed); + } else if (key == "temporal_decode_frames") { + config.decode_frames = std::max(1, parsed); } else { config.tile_frames = std::max(1, parsed); } diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 9a05839cfd..5285a0d738 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -1288,9 +1288,11 @@ namespace WAN { ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx, ggml_tensor* z, int chunk_idx, - int64_t b = 1) { + int64_t b = 1, + int decode_frames = 1) { // z: [b*c, t, h, w] GGML_ASSERT(b == 1); + GGML_ASSERT(decode_frames >= 1); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); @@ -1304,11 +1306,13 @@ namespace WAN { } ggml_tensor* out = nullptr; - for (int64_t frame = 0; frame < x->ne[2]; ++frame) { + for (int64_t frame = 0; frame < x->ne[2];) { const int global_frame = chunk_idx + static_cast(frame); - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1); - _conv_idx = 0; - auto out_frame = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame); + // The first latent frame bypasses temporal upsampling and primes the causal caches. + const int64_t count = !is_2D && global_frame > 0 ? std::min(decode_frames, x->ne[2] - frame) : 1; + auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + count); + _conv_idx = 0; + auto out_frame = decoder->forward(ctx, in, b, _feat_map, _conv_idx, global_frame); if (is_2D && global_frame > 0) { auto repeated = out_frame; for (int repeat = 1; repeat < 4; ++repeat) { @@ -1317,6 +1321,7 @@ namespace WAN { out_frame = repeated; } out = out == nullptr ? out_frame : ggml_concat(ctx->ggml_ctx, out, out_frame, 2); + frame += count; } return unpatchify(ctx->ggml_ctx, out, patch_size, b); } @@ -1456,7 +1461,7 @@ namespace WAN { return gf; } - ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx) { + ggml_cgraph* build_temporal_tile_graph(const sd::Tensor& z_tensor, int chunk_idx, int decode_frames) { ggml_cgraph* gf = new_graph_custom(std::max(20480, 10240 * z_tensor.shape()[2])); ae.clear_cache(); @@ -1470,7 +1475,8 @@ namespace WAN { auto runner_ctx = get_context(gf); - ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx); + ggml_tensor* out = ae.decode_tiled_chunk(&runner_ctx, z, chunk_idx, 1, + version == VERSION_WAN2_2_TI2V ? decode_frames : 1); for (size_t feat_idx = 0; feat_idx < ae._feat_map.size(); feat_idx++) { ggml_tensor* feat_cache = ae._feat_map[feat_idx]; @@ -1493,8 +1499,19 @@ namespace WAN { stateful_config.overlap = 0; auto plan = make_vae_temporal_tile_plan(input.shape()[2], stateful_config); - LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, total latent frames=%lld, tiles=%d", + if (sd_version_is_wan(version) && !ae.is_2D && plan.tiles.size() > 1 && + plan.tiles.size() <= 3 && plan.tile_frames > 2 && + plan.tiles.back().end - plan.tiles.back().start == 1) { + // Longer plans retain full middle tiles, so balancing their tail cannot shrink the largest decoded output. + auto& previous = plan.tiles[plan.tiles.size() - 2]; + auto& last = plan.tiles.back(); + --previous.end; + --last.start; + } + + LOG_VERBOSE("Wan VAE stateful temporal tiling: tile_frames=%d, decode_frames=%d, total latent frames=%lld, tiles=%d", plan.tile_frames, + version == VERSION_WAN2_2_TI2V && !ae.is_2D ? config.decode_frames : 1, (long long)input.shape()[2], (int)plan.tiles.size()); @@ -1508,7 +1525,7 @@ namespace WAN { (long long)tile.start, (long long)tile.end); auto get_graph = [&]() -> ggml_cgraph* { - return build_temporal_tile_graph(input_tile, static_cast(tile.start)); + return build_temporal_tile_graph(input_tile, static_cast(tile.start), config.decode_frames); }; return restore_trailing_singleton_dims( GGMLRunner::compute(get_graph, n_threads, false), diff --git a/src/model_definition.h b/src/model_definition.h new file mode 100644 index 0000000000..2b49994958 --- /dev/null +++ b/src/model_definition.h @@ -0,0 +1,47 @@ +#ifndef __SD_MODEL_DEFINITION_H__ +#define __SD_MODEL_DEFINITION_H__ + +#include +#include "model.h" +#include "stable-diffusion.h" + +class TokenizerConfig; +struct VAE; +namespace sd::model_builders { + struct Context; + struct CoreRunners; + struct VAEOptions; +} + +struct ModelSamplingPolicy { + prediction_t default_prediction = EPS_PRED; + float default_flow_shift = std::numeric_limits::infinity(); + sample_method_t default_sample_method = EULER_A_SAMPLE_METHOD; + scheduler_t default_scheduler = DISCRETE_SCHEDULER; + + bool probe_v_prediction = false; + + prediction_t (*prediction_from_weights)(const String2TensorStorage&) = nullptr; + float (*flow_shift_from_weights)(const String2TensorStorage&) = nullptr; +}; + +struct ModelGeometry { + int diffusion_down_factor = 8; + int latent_channels = 4; + int image_channels = 3; +}; + +struct ModelDefinition { + SDVersion version = VERSION_COUNT; + const char* name = nullptr; + ModelTraits traits; + ModelSamplingPolicy sampling; + ModelGeometry geometry; + const char* ref_image_preset = "default"; + void (*configure_vae)(const sd::model_builders::Context&, const sd::model_builders::VAEOptions&, VAE&) = nullptr; + bool (*build)(const sd::model_builders::Context&, const TokenizerConfig&, sd::model_builders::CoreRunners&) = nullptr; +}; + +const ModelDefinition* find_model_definition(SDVersion version); + +#endif // __SD_MODEL_DEFINITION_H__ diff --git a/src/model_io/gguf_reader_ext.h b/src/model_io/gguf_reader_ext.h index f02fe3c6f4..bf183f7488 100644 --- a/src/model_io/gguf_reader_ext.h +++ b/src/model_io/gguf_reader_ext.h @@ -60,7 +60,12 @@ class GGUFReader { if (count > remaining_bytes_ / element_size) return false; uint64_t size = count * element_size; - fin.seekg(static_cast(size), std::ios::cur); + // Seeking discards the read buffer, making arrays of short tokenizer strings expensive to skip. + if (size <= 64 * 1024) { + fin.ignore(static_cast(size)); + } else { + fin.seekg(static_cast(size), std::ios::cur); + } if (!fin.good()) return false; remaining_bytes_ -= size; diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 90ab65eed5..3bfd9e55ca 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -132,6 +132,10 @@ void convert_tensor(void* src, } else if (dst_type == GGML_TYPE_F32) { if (src_type == GGML_TYPE_F16) { ggml_fp16_to_fp32_row((ggml_fp16_t*)src, (float*)dst, n); + } else if (src_type == GGML_TYPE_I32) { + for (int i = 0; i < n; i++) { + ((float*)dst)[i] = static_cast(((int32_t*)src)[i]); + } } else { auto qtype = ggml_get_type_traits(src_type); if (qtype->to_float == nullptr) { @@ -340,7 +344,7 @@ bool ModelLoader::init_from_safetensors_index_file(const std::string& file_path, const auto expected = (std::filesystem::u8path(file_path).parent_path() / std::filesystem::u8path(entry.value().get())) .lexically_normal(); - bool found = false; + bool found = false; for (const auto& shard : parsed_tensor_names_) { if (std::filesystem::u8path(shard.first).lexically_normal() == expected) { found = shard.second.count(entry.key()) != 0; @@ -439,6 +443,10 @@ bool ModelLoader::init_from_diffusers_file(const std::string& file_path, const s SDVersion ModelLoader::get_sd_version() const { TensorStorage token_embedding_weight, input_block_weight, context_ebedding_weight; + if (tensor_storage_map.count("model.diffusion_model.perceiver.latent_queries")) { + return VERSION_FIBO2; + } + if (tensor_storage_map.find("model.diffusion_model.emb_in.image.weight") != tensor_storage_map.end() && tensor_storage_map.find("model.diffusion_model.txt_mode_blocks.0.q_proj.weight") != tensor_storage_map.end()) { return VERSION_FLUX3; @@ -464,6 +472,8 @@ SDVersion ModelLoader::get_sd_version() const { bool has_output_block_311 = false; bool has_output_block_71 = false; bool has_attn_1024 = false; + bool has_distilled_guidance = false; + bool has_sdxl_add_embedding = false; for (auto& [name, tensor_storage] : tensor_storage_map) { if (tensor_storage.name.find("model.diffusion_model.double_blocks.") != std::string::npos || @@ -483,6 +493,9 @@ SDVersion ModelLoader::get_sd_version() const { if (tensor_storage.name.find("model.diffusion_model.nerf_final_layer_conv.") != std::string::npos) { return VERSION_CHROMA_RADIANCE; } + if (tensor_storage.name.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { + has_distilled_guidance = true; + } if (tensor_storage.name.find("model.diffusion_model.joint_blocks.") != std::string::npos) { return VERSION_SD3; } @@ -565,6 +578,25 @@ SDVersion ModelLoader::get_sd_version() const { } if (tensor_storage.name.find("model.diffusion_model.video_patch_proj.weight") != std::string::npos && tensor_storage_map.find("model.diffusion_model.audio_patch_proj.weight") != tensor_storage_map.end()) { + // P-H3 uses the same DiT architecture as H3; its VAE or connector identifies the image pipeline. + auto decoder = tensor_storage_map.find("first_stage_model.decoder.proj_out.weight"); + if (decoder == tensor_storage_map.end()) { + decoder = tensor_storage_map.find("vae.decoder.proj_out.weight"); + } + if (decoder != tensor_storage_map.end() && decoder->second.ne[1] == 768) { + return VERSION_P_H3_IMAGE; + } + for (const std::string prefix : {"", "embeddings_connectors."}) { + if (tensor_storage_map.count(prefix + "text_encoders.ph3_image_connector.proj.weight") || + tensor_storage_map.count(prefix + "text_encoders.llm.text_connector.proj.weight")) { + return VERSION_P_H3_IMAGE; + } + } + auto connector = tensor_storage_map.find("embeddings_connectors.proj.weight"); + if (connector != tensor_storage_map.end() && connector->second.ne[0] == 12800 && + connector->second.ne[1] == 5120 && tensor_storage_map.count("embeddings_connectors.sink_embedding")) { + return VERSION_P_H3_IMAGE; + } return VERSION_MINIMAX_H3; } if (name == "model.diffusion_model.y_embedder.layer_pool.weight" && @@ -589,6 +621,7 @@ SDVersion ModelLoader::get_sd_version() const { has_img_emb = true; } if (tensor_storage.name.find("model.diffusion_model.input_blocks.") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.down_blocks.") != std::string::npos || tensor_storage.name.find("unet.down_blocks.") != std::string::npos) { is_unet = true; if (has_multiple_encoders) { @@ -603,18 +636,27 @@ SDVersion ModelLoader::get_sd_version() const { is_xl = true; } } + // A standalone SDXL UNet (--diffusion-model) identifies the family by its added + // text/time embedding; its clip_g never carries a checkpoint prefix. + if (tensor_storage.name.find("model.diffusion_model.add_embedding.") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.label_emb.0.0.") != std::string::npos) { + has_sdxl_add_embedding = true; + } if (tensor_storage.name.find("model.diffusion_model.input_blocks.8.0.time_mixer.mix_factor") != std::string::npos) { return VERSION_SVD; } if (tensor_storage.name.find("model.diffusion_model.middle_block.1.") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.mid_block.resnets.1.") != std::string::npos || tensor_storage.name.find("unet.mid_block.resnets.1.") != std::string::npos) { has_middle_block_1 = true; } if (tensor_storage.name.find("model.diffusion_model.output_blocks.3.1.transformer_blocks.1") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.up_blocks.1.attentions.0.transformer_blocks.1") != std::string::npos || tensor_storage.name.find("unet.up_blocks.1.attentions.0.transformer_blocks.1") != std::string::npos) { has_output_block_311 = true; } if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1") != std::string::npos || + tensor_storage.name.find("model.diffusion_model.up_blocks.2.attentions.1") != std::string::npos || tensor_storage.name.find("unet.up_blocks.2.attentions.1") != std::string::npos) { has_output_block_71 = true; if (tensor_storage.name.find("model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight") != std::string::npos) { @@ -632,6 +674,7 @@ SDVersion ModelLoader::get_sd_version() const { // break; } if (tensor_storage.name == "model.diffusion_model.input_blocks.0.0.weight" || + tensor_storage.name == "model.diffusion_model.conv_in.weight" || tensor_storage.name == "model.diffusion_model.img_in.weight" || tensor_storage.name == "unet.conv_in.weight") { input_block_weight = tensor_storage; @@ -655,6 +698,9 @@ SDVersion ModelLoader::get_sd_version() const { } bool is_inpaint = input_block_weight.ne[2] == 9; bool is_ip2p = input_block_weight.ne[2] == 8; + if (is_unet && has_sdxl_add_embedding) { + is_xl = true; + } if (is_xl) { if (is_inpaint) { return VERSION_SDXL_INPAINT; @@ -672,6 +718,9 @@ SDVersion ModelLoader::get_sd_version() const { } if (is_flux && !is_flux2) { + if (has_distilled_guidance) { + return VERSION_CHROMA; + } if (context_ebedding_weight.ne[0] == 3584) { return VERSION_LONGCAT; } else { diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 05ee55f676..28a1636dc8 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1,10 +1,12 @@ #include "model_manager.h" #include +#include #include #include #include #include +#include #include #include @@ -571,7 +573,7 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vector& states // The temporary runner is destroyed before this manager call returns. auto borrowed_manager = std::shared_ptr(this, [](ModelManager*) {}); auto lora = std::make_shared(id, compute_backend, target->params_backend, - borrowed_manager, lora_spec.file_id, lora_version_, - target->residency_mode); + borrowed_manager, lora_spec.file_id, lora_version_, + target->residency_mode); LoraModel::filter_t lora_tensor_filter = nullptr; if (!lora_spec.tensor_name_prefix_filter.empty()) { @@ -1466,8 +1468,8 @@ bool ModelManager::assign_compute_backend(const std::vector& tenso const bool params_follow_compute = state->params_follow_compute_backend || state->residency_mode == ResidencyMode::Disk; - const bool compute_changes = state->compute_backend != compute_backend; - const bool params_changes = params_follow_compute && state->params_backend != compute_backend; + const bool compute_changes = state->compute_backend != compute_backend; + const bool params_changes = params_follow_compute && state->params_backend != compute_backend; if (!compute_changes && !params_changes) { continue; } @@ -1559,7 +1561,7 @@ size_t ModelManager::compute_backend_resident_bytes(ggml_backend_t compute_backe ? buffer_type != nullptr && ggml_backend_buft_get_device(buffer_type) == compute_device : std::any_of(split_devices->second.begin(), split_devices->second.end(), [&](const auto& entry) { return ggml_backend_get_device(entry.first) == compute_device; - }); + }); if (!on_device) { return; } @@ -1828,10 +1830,16 @@ ModelManager::CapacityCheck ModelManager::check_capacity( auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; }; const size_t missing = compute_backend_alloc_size(states, true); // Backend scratch buffers and pipelines are not included in graph measurements. - constexpr size_t safety_margin = 512ULL * 1024ULL * 1024ULL; - result.required_device_bytes = add(add(request.pending_allocation_bytes, missing), safety_margin); - result.required_budget_bytes = add(request.runtime_peak_bytes(), missing); - auto available_device_bytes = [&](ggml_backend_t backend) { + size_t safety_margin = 512ULL * 1024ULL * 1024ULL; +#ifdef _WIN32 + if (sd_backend_is(request.compute_backend, "CUDA")) { + // Avoid WDDM allocation churn when transient buffers compete with reusable staging buffers. + safety_margin = 1024ULL * 1024ULL * 1024ULL; + } +#endif + result.required_device_bytes = add(add(request.pending_allocation_bytes, missing), safety_margin); + result.required_budget_bytes = add(request.runtime_peak_bytes(), missing); + auto available_device_bytes = [&](ggml_backend_t backend) { auto device = ggml_backend_get_device(backend); if (device == nullptr) { return SIZE_MAX; @@ -1857,10 +1865,10 @@ ModelManager::CapacityCheck ModelManager::check_capacity( const size_t resident = add(weights_resident, add(other_runtime, request.runtime_resident_bytes)); if (log_details) { LOG_WARN("model manager memory on %s: reported free %.2f MB / total %.2f MB, tracked weights %.2f MB / other runtime %.2f MB / current runtime %.2f MB", - ggml_backend_name(backend), - free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0), - weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0), - request.runtime_resident_bytes / (1024.0 * 1024.0)); + ggml_backend_name(backend), + free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0), + weights_resident / (1024.0 * 1024.0), other_runtime / (1024.0 * 1024.0), + request.runtime_resident_bytes / (1024.0 * 1024.0)); } if (free_bytes == 0 && total_bytes == 0) { return SIZE_MAX; @@ -1872,7 +1880,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( } // A wrapped report may reflect real pressure; reclaim before accepting the tracked estimate. LOG_VERBOSE("ignoring impossible free memory report on %s (%.2f MB free / %.2f MB total), using tracked residency", - ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0)); + ggml_backend_name(backend), free_bytes / (1024.0 * 1024.0), total_bytes / (1024.0 * 1024.0)); free_bytes = total_bytes; result.requires_reclamation = true; } @@ -1995,21 +2003,21 @@ bool ModelManager::ensure_compute_backend_capacity( } const bool reloadable = state->residency_mode == ResidencyMode::Disk || state->compute_backend != state->params_backend; - const bool resident = state->compute_backend == state->params_backend - ? state->loaded_to_params_backend - : state->staged_to_compute_backend; + const bool resident = state->compute_backend == state->params_backend + ? state->loaded_to_params_backend + : state->staged_to_compute_backend; if (reloadable && resident) { eviction_states.insert(state); } }; size_t last_eviction_count = 0; auto release_eviction_states = [&]() { - if (eviction_states.size() == last_eviction_count) { - return false; + if (eviction_states.size() != last_eviction_count) { + last_eviction_count = eviction_states.size(); + release_compute_staging_blocks(false, &eviction_states); + release_params_storage_blocks(false, &eviction_states); } - last_eviction_count = eviction_states.size(); - release_compute_staging_blocks(false, &eviction_states); - release_params_storage_blocks(false, &eviction_states); + // Device memory reports can lag behind releases even when no new candidates are added. return fits(); }; @@ -2088,10 +2096,30 @@ bool ModelManager::ensure_compute_backend_capacity( } revert_demotions(); - if (check_capacity(request, required_states).fits()) { + auto capacity = check_capacity(request, required_states); + if (!capacity.requires_reclamation && capacity.fits()) { return true; } - const auto capacity = check_capacity(request, required_states, true); + if (sd_backend_is(compute_backend, "Vulkan") && + capacity.required_budget_bytes <= capacity.available_budget_bytes) { + ggml_backend_synchronize(compute_backend); + sd_backend_clear_buffer_cache(compute_backend); + // Vulkan heap usage can lag behind buffer releases. Fast memory queries + // must not turn that reporting delay into an allocation failure. + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100); + for (;;) { + capacity = check_capacity(request, required_states); + if (!capacity.requires_reclamation && capacity.fits()) { + return true; + } + if (capacity.required_budget_bytes > capacity.available_budget_bytes || + std::chrono::steady_clock::now() >= deadline) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } + } + capacity = check_capacity(request, required_states, true); const std::string available_device = capacity.available_device_bytes == SIZE_MAX ? "unknown" : sd_format("%.2f MB", capacity.available_device_bytes / (1024.0 * 1024.0)); diff --git a/src/model_manager.h b/src/model_manager.h index f6d442a8d5..3f21a4be38 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -77,6 +77,7 @@ class ModelManager : public DeviceResidencyManager { ggml_backend_t compute_backend = nullptr; ggml_backend_t transfer_backend = nullptr; ggml_backend_event_t event = nullptr; + bool submitted = false; std::vector> staging_blocks; }; @@ -298,6 +299,10 @@ class ModelManager : public DeviceResidencyManager { WeightPrefetchResult prefetch_params( uintptr_t owner_id, const std::vector& tensors) override; + WeightPrefetchResult prepare_prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) override; + void submit_prefetched_params(uintptr_t owner_id) override; bool activate_prefetched_params(uintptr_t owner_id, const std::vector& tensors) override; void clear_prefetched_params(uintptr_t owner_id) override; diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp index de96d5ae3b..4056b8f0ea 100644 --- a/src/model_manager_prefetch.cpp +++ b/src/model_manager_prefetch.cpp @@ -33,7 +33,7 @@ void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) { ggml_backend_event_synchronize(block.event); ggml_backend_event_free(block.event); block.event = nullptr; - } else if (block.transfer_backend != nullptr) { + } else if (block.submitted && block.transfer_backend != nullptr) { ggml_backend_synchronize(block.transfer_backend); } block.transfer_backend = nullptr; @@ -89,7 +89,7 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { backend_limit = MAX_RESIDENCY_BLOCK_BYTES; } - auto enqueue_chunk = [&](const std::vector& chunk) -> bool { + auto allocate_chunk = [&](const std::vector& chunk) -> bool { if (chunk.empty()) { return true; } @@ -128,23 +128,6 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { ggml_backend_buffer_set_usage(staging_block->buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - for (const auto& pair : staging_block->staged_tensors) { - TensorState* state = pair.first; - ggml_tensor* staging_tensor = pair.second; - const bool host_source = ggml_backend_buffer_is_host(state->tensor->buffer); - if (host_source) { - ggml_backend_tensor_set_async(block.transfer_backend, - staging_tensor, - state->tensor->data, - 0, - ggml_nbytes(state->tensor)); - } else { - ggml_backend_tensor_copy_async(state->params_backend, - block.transfer_backend, - state->tensor, - staging_tensor); - } - } block.staging_blocks.push_back(std::move(staging_block)); return true; }; @@ -156,7 +139,7 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { ggml_backend_buft_get_alloc_size(buffer_type, state->tensor), alignment); if (!chunk.empty() && backend_limit > 0 && tensor_size > backend_limit - std::min(chunk_size, backend_limit)) { - if (!enqueue_chunk(chunk)) { + if (!allocate_chunk(chunk)) { return false; } chunk.clear(); @@ -165,10 +148,39 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { chunk.push_back(state); chunk_size = tensor_size > SIZE_MAX - chunk_size ? SIZE_MAX : chunk_size + tensor_size; } - if (!enqueue_chunk(chunk)) { + if (!allocate_chunk(chunk)) { return false; } + return true; +} + +void ModelManager::submit_prefetched_params(uintptr_t owner_id) { + auto existing = prefetch_blocks_.find(owner_id); + if (existing == prefetch_blocks_.end() || existing->second->submitted) { + return; + } + auto& block = *existing->second; + block.submitted = true; + for (const auto& staging_block : block.staging_blocks) { + for (const auto& pair : staging_block->staged_tensors) { + TensorState* state = pair.first; + ggml_tensor* staging_tensor = pair.second; + if (ggml_backend_buffer_is_host(state->tensor->buffer)) { + ggml_backend_tensor_set_async(block.transfer_backend, + staging_tensor, + state->tensor->data, + 0, + ggml_nbytes(state->tensor)); + } else { + ggml_backend_tensor_copy_async(state->params_backend, + block.transfer_backend, + state->tensor, + staging_tensor); + } + } + } + ggml_backend_dev_t device = ggml_backend_get_device(block.transfer_backend); block.event = ggml_backend_event_new(device); if (block.event != nullptr) { @@ -186,10 +198,19 @@ bool ModelManager::populate_prefetch_block(PrefetchBlock& block) { total_size / (1024.f * 1024.f), block.states.size(), ggml_backend_name(block.compute_backend)); - return true; } WeightPrefetchResult ModelManager::prefetch_params( + uintptr_t owner_id, + const std::vector& tensors) { + auto result = prepare_prefetch_params(owner_id, tensors); + if (result == WeightPrefetchResult::Scheduled) { + submit_prefetched_params(owner_id); + } + return result; +} + +WeightPrefetchResult ModelManager::prepare_prefetch_params( uintptr_t owner_id, const std::vector& tensors) { if (tensors.empty()) { @@ -288,6 +309,7 @@ bool ModelManager::activate_prefetched_params( if (existing == prefetch_blocks_.end()) { return false; } + submit_prefetched_params(owner_id); std::unique_ptr block = std::move(existing->second); prefetch_blocks_.erase(existing); synchronize_prefetch_block(*block); diff --git a/src/model_registry.cpp b/src/model_registry.cpp new file mode 100644 index 0000000000..9a80bb8359 --- /dev/null +++ b/src/model_registry.cpp @@ -0,0 +1,102 @@ +#include +#include +#include "model_definition.h" + +const std::vector& fibo2_definitions(); +const std::vector& sd3_definitions(); +const std::vector& pid_definitions(); +const std::vector& ideogram4_definitions(); +const std::vector& krea2_definitions(); +const std::vector& flux_definitions(); +const std::vector& flux3_definitions(); +const std::vector& flux2_definitions(); +const std::vector& ltxav_definitions(); +const std::vector& minimax_h3_definitions(); +const std::vector& hunyuan_video_definitions(); +const std::vector& wan_definitions(); +const std::vector& lingbot_video_definitions(); +const std::vector& qwen_image_definitions(); +const std::vector& pixart_definitions(); +const std::vector& mage_flow_definitions(); +const std::vector& longcat_definitions(); +const std::vector& hidream_o1_definitions(); +const std::vector& minit2i_definitions(); +const std::vector& sensenova_u1_definitions(); +const std::vector& anima_definitions(); +const std::vector& ming_image_definitions(); +const std::vector& z_image_definitions(); +const std::vector& z_image_l2p_definitions(); +const std::vector& llada_image_definitions(); +const std::vector& boogu_image_definitions(); +const std::vector& ernie_image_definitions(); +const std::vector& lens_definitions(); +const std::vector& iris_definitions(); +const std::vector& sd_definitions(); +const std::vector& moss_image_definitions(); + +static const std::array& model_definitions() { + static const auto definitions = [] { + std::array result{}; + auto append = [&](const std::vector& family) { + for (const auto& definition : family) { + if (definition.version < 0 || definition.version >= VERSION_COUNT || result[definition.version]) { + throw std::logic_error("duplicate or invalid model registration"); + } + result[definition.version] = &definition; + } + }; + append(fibo2_definitions()); + append(sd3_definitions()); + append(pid_definitions()); + append(ideogram4_definitions()); + append(krea2_definitions()); + append(flux_definitions()); + append(flux3_definitions()); + append(flux2_definitions()); + append(ltxav_definitions()); + append(minimax_h3_definitions()); + append(hunyuan_video_definitions()); + append(wan_definitions()); + append(lingbot_video_definitions()); + append(qwen_image_definitions()); + append(pixart_definitions()); + append(mage_flow_definitions()); + append(longcat_definitions()); + append(hidream_o1_definitions()); + append(minit2i_definitions()); + append(sensenova_u1_definitions()); + append(anima_definitions()); + append(ming_image_definitions()); + append(z_image_definitions()); + append(z_image_l2p_definitions()); + append(llada_image_definitions()); + append(boogu_image_definitions()); + append(ernie_image_definitions()); + append(lens_definitions()); + append(iris_definitions()); + append(sd_definitions()); + append(moss_image_definitions()); + for (const auto* definition : result) { + if (!definition) + throw std::logic_error("missing model registration"); + } + return result; + }(); + return definitions; +} + +const ModelDefinition* find_model_definition(SDVersion version) { + if (version < 0 || version >= VERSION_COUNT) + return nullptr; + return model_definitions()[version]; +} + +const ModelTraits* find_model_traits(SDVersion version) { + const auto* definition = find_model_definition(version); + return definition ? &definition->traits : nullptr; +} + +const char* sd_version_name(SDVersion version) { + const auto* definition = find_model_definition(version); + return definition ? definition->name : "Unknown"; +} diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 6a519ba21f..2e147348ff 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -197,7 +197,11 @@ std::string convert_cond_stage_model_name(std::string name, std::string prefix, replace_with_name_map(name, llm_vision_name_map); } else { replace_with_name_map(name, llm_name_map); - name = convert_llada2_moe_te_name(name); + name = convert_llada2_moe_te_name(name); + const std::string language_prefix = "llm.model.language_model."; + if (starts_with(name, language_prefix)) { + name = "llm.model." + name.substr(language_prefix.size()); + } } } else { name = convert_open_clip_to_hf_clip_name(name); @@ -704,6 +708,24 @@ std::string convert_diffusers_dit_to_original_flux(std::string name) { return name; } +static std::string convert_diffusers_dit_to_original_flux3(std::string name) { + static const std::vector> prefix_map = { + {"single_transformer_blocks.", "single_blocks."}, + {"mode_blocks.", "content_mode_blocks."}, + }; + static const std::vector> name_map = { + {".attn.to_q.", ".q_proj."}, + {".attn.to_k.", ".k_proj."}, + {".attn.to_v.", ".v_proj."}, + {".attn.to_out.0.", ".attn_out."}, + {".ff.linear_in.", ".mlp_in."}, + {".ff.linear_out.", ".mlp_out."}, + }; + replace_with_prefix_map(name, prefix_map); + replace_with_name_map(name, name_map); + return name; +} + std::string convert_hunyuan_video_to_original_flux(std::string name) { int num_layers = 54; int num_single_layers = 0; @@ -1034,6 +1056,18 @@ static std::string convert_ming_image_dit_name(std::string name) { return name; } +static std::string convert_fibo2_dit_name(std::string name) { + static const std::vector> name_map = { + {".attention.q_norm.", ".attention.norm_q."}, + {".attention.k_norm.", ".attention.norm_k."}, + {".attention.out.", ".attention.to_out.0."}, + {".cross_attention.to_out.weight", ".cross_attention.to_out.0.weight"}, + {".self_attention.to_out.weight", ".self_attention.to_out.0.weight"}, + }; + replace_with_name_map(name, name_map); + return name; +} + std::string convert_diffusers_dit_to_original_minimax_h3(std::string name) { static const std::vector> prefix_map = { {"transformer_blocks.", "blocks."}, @@ -1061,12 +1095,16 @@ std::string convert_diffusion_model_name(std::string name, std::string prefix, S name = convert_diffusers_dit_to_original_sd3(name); } else if (sd_version_is_flux(version) || sd_version_is_flux2(version) || sd_version_is_longcat(version) || sd_version_is_sefi_image(version)) { name = convert_diffusers_dit_to_original_flux(name); + } else if (sd_version_is_flux3(version)) { + name = convert_diffusers_dit_to_original_flux3(name); } else if (sd_version_is_hunyuan_video(version)) { name = convert_hunyuan_video_to_original_flux(name); } else if (version == VERSION_MOSS_IMAGE) { name = convert_moss_image_dit_name(name); } else if (version == VERSION_MING_IMAGE) { name = convert_ming_image_dit_name(name); + } else if (version == VERSION_FIBO2) { + name = convert_fibo2_dit_name(name); } else if (sd_version_is_z_image(version) || sd_version_is_z_image_l2p(version)) { name = convert_diffusers_dit_to_original_lumina2(name); } else if (sd_version_is_llada_image(version)) { @@ -1661,12 +1699,19 @@ std::string convert_tensor_name(std::string name, SDVersion version) { } } + const std::vector> connector_prefix_map = { + {"embeddings_connectors.text_encoders.", "text_encoders."}, + {"embeddings_connectors.", version == VERSION_P_H3_IMAGE ? "text_encoders.ph3_image_connector." : ""}, + }; + replace_with_prefix_map(name, connector_prefix_map); + std::unordered_map prefix_map = { {"diffusion_model.", "model.diffusion_model."}, {"unet.", "model.diffusion_model."}, {"transformer.", "model.diffusion_model."}, // dit {"vae.", "first_stage_model."}, {"text_encoders.llm.text_embedding_projection.", "text_embedding_projection."}, + {"text_encoders.llm.text_connector.", "text_encoders.ph3_image_connector."}, {"text_encoder.", "cond_stage_model.transformer."}, {"te.", "cond_stage_model.transformer."}, {"text_encoder.2.", "cond_stage_model.1.transformer."}, @@ -1715,7 +1760,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { } } - if (version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) { + if (version == VERSION_FIBO2 || version == VERSION_QWEN_IMAGE_2_1 || sd_version_is_boogu_image(version) || sd_version_is_krea2(version) || sd_version_is_mage_flow(version) || sd_version_is_minimax_h3(version)) { const std::string hf_vision_prefix = "text_encoders.llm.model.visual."; if (starts_with(name, hf_vision_prefix)) { name = "text_encoders.llm.visual." + name.substr(hf_vision_prefix.size()); diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 82b40473a5..a4a6d9ddf8 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -22,6 +22,7 @@ #include "core/probe/conv_direct.h" #include "core/probe/flash_attn.h" #include "model.h" +#include "model_definition.h" #include "core/rng.hpp" #include "core/rng_mt19937.hpp" @@ -42,6 +43,7 @@ #include "model/diffusion/animatediff.hpp" #include "model/diffusion/control.hpp" #include "model/diffusion/model.hpp" +#include "model/diffusion/z_image.hpp" #include "model/vae/audio_vae.hpp" #include "model/vae/ltx_vae.hpp" #include "model/vae/vae.hpp" @@ -56,69 +58,6 @@ #include -const char* model_version_to_str[] = { - "SD 1.x", - "SD 1.x Inpaint", - "Instruct-Pix2Pix", - "SD 1.x Tiny UNet", - "SD 2.x", - "SD 2.x Inpaint", - "SD 2.x Tiny UNet", - "SDXS (512-DS)", - "SDXS (09)", - "SDXL", - "SDXL Inpaint", - "SDXL Instruct-Pix2Pix", - "SDXL (Vega)", - "SDXL (SSD1B)", - "SVD", - "SD3.x", - "Flux", - "Flux Fill", - "Flux Control", - "Flex.2", - "Chroma Radiance", - "Wan 2.x", - "Wan 2.2 I2V", - "Wan 2.2 TI2V", - "Wan 2.2 S2V", - "LingBot Video", - "Qwen Image", - "Qwen Image Layered", - "Qwen Image 2.1", - "Hunyuan Video", - "Anima", - "Flux.2", - "Flux.2 klein", - "LTXAV", - "MiniMax-H3", - "HiDream O1", - "Z-Image", - "Boogu Image", - "Ovis Image", - "Ernie Image", - "Lens", - "MiniT2I", - "Longcat-Image", - "PiD", - "Ideogram 4", - "SeFi-Image", - "Krea2", - "Mage Flow", - "SenseNova U1.5", - "LLaDA-Image", - "ESRGAN", - "PixArt", - "Ming-Image", - "Z-Image L2P", - "Iris-3B", - "Flux.3 Image Klein", - "MOSS-Image", -}; - -static_assert(VERSION_COUNT == sizeof(model_version_to_str) / sizeof(model_version_to_str[0]), - "\nnumber of elements in model_version_to_str[] != VERSION_COUNT"); - void calculate_alphas_cumprod(float* alphas_cumprod, float linear_start = 0.00085f, float linear_end = 0.0120f, @@ -297,6 +236,15 @@ bool StableDiffusionGGML::apply_model_update(ModelLoader candidate, } } } + if (sd_version_is_z_image(version) && sd_version_is_z_image(next_version)) { + if (groups.count(RunnerGroup::ControlNet) > 0) { + groups.insert(RunnerGroup::Core); + } + if (groups.count(RunnerGroup::Core) > 0 && + model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) { + groups.insert(RunnerGroup::ControlNet); + } + } runner_state_.ready = false; if (!reset_runners(groups)) { return false; @@ -470,7 +418,7 @@ bool StableDiffusionGGML::register_row_split_runner_params(ModelComponent compon size_t free_bytes = 0, total_bytes = 0; ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); int64_t usable_bytes = std::max((int64_t)free_bytes - compute_headroom_bytes, - (int64_t)free_bytes / 8); + (int64_t)free_bytes / 8); tensor_split[reg_index] = usable_bytes > 0 ? (float)((double)usable_bytes / (1024.0 * 1024.0)) : 1.0f; } @@ -574,8 +522,8 @@ bool StableDiffusionGGML::register_layer_split_runner_params(ModelComponent comp model->set_runtime_backends(module_backends); model->set_graph_cut_layer_split_backend_vram_limits(layer_split_vram_limits_for_backends(module_backends)); model->set_graph_cut_layer_split_enabled(true); - const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) || - backend_manager.params_backend_is_disk(module); + const bool params_follow_runtime = backend_manager.params_backend_follows_runtime(module) || + backend_manager.params_backend_is_disk(module); ggml_backend_t initial_params_backend = params_follow_runtime ? module_backends[0] : params_backend_for(module); if (initial_params_backend == nullptr) { return false; @@ -728,87 +676,93 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf const auto* sd_ctx_params = &configuration.params; auto& use_tae = configuration.use_tae; auto& use_audio_vae = configuration.use_audio_vae; + const auto load_model = [&](const char* name, const char* path, const std::string& prefix = "") { + if (!model_loader.init_from_file(path, prefix)) + return false; + configuration.record_model(name, path); + return true; + }; if (strlen(SAFE_STR(sd_ctx_params->model_path)) > 0) { LOG_INFO("loading model from '%s'", sd_ctx_params->model_path); - if (!model_loader.init_from_file(sd_ctx_params->model_path)) { + if (!load_model("model", sd_ctx_params->model_path)) { LOG_ERROR("init model loader from file failed: '%s'", sd_ctx_params->model_path); } } if (strlen(SAFE_STR(sd_ctx_params->diffusion_model_path)) > 0) { LOG_INFO("loading diffusion model from '%s'", sd_ctx_params->diffusion_model_path); - if (!model_loader.init_from_file(sd_ctx_params->diffusion_model_path, "model.diffusion_model.")) { + if (!load_model("diffusion_model", sd_ctx_params->diffusion_model_path, "model.diffusion_model.")) { LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->diffusion_model_path); } } if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { LOG_INFO("loading high noise diffusion model from '%s'", sd_ctx_params->high_noise_diffusion_model_path); - if (!model_loader.init_from_file(sd_ctx_params->high_noise_diffusion_model_path, "model.high_noise_diffusion_model.")) { + if (!load_model("high_noise_diffusion_model", sd_ctx_params->high_noise_diffusion_model_path, "model.high_noise_diffusion_model.")) { LOG_WARN("loading diffusion model from '%s' failed", sd_ctx_params->high_noise_diffusion_model_path); } } if (strlen(SAFE_STR(sd_ctx_params->uncond_diffusion_model_path)) > 0) { LOG_INFO("loading unconditional diffusion model from '%s'", sd_ctx_params->uncond_diffusion_model_path); - if (!model_loader.init_from_file(sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) { + if (!load_model("uncond_diffusion_model", sd_ctx_params->uncond_diffusion_model_path, "model.diffusion_model.uncond.")) { LOG_WARN("loading unconditional diffusion model from '%s' failed", sd_ctx_params->uncond_diffusion_model_path); } } if (strlen(SAFE_STR(sd_ctx_params->clip_l_path)) > 0) { LOG_INFO("loading clip_l from '%s'", sd_ctx_params->clip_l_path); - if (!model_loader.init_from_file(sd_ctx_params->clip_l_path, "clip_l.")) { + if (!load_model("clip_l", sd_ctx_params->clip_l_path, "clip_l.")) { LOG_WARN("loading clip_l from '%s' failed", sd_ctx_params->clip_l_path); } } if (strlen(SAFE_STR(sd_ctx_params->clip_g_path)) > 0) { LOG_INFO("loading clip_g from '%s'", sd_ctx_params->clip_g_path); - if (!model_loader.init_from_file(sd_ctx_params->clip_g_path, "clip_g.")) { + if (!load_model("clip_g", sd_ctx_params->clip_g_path, "clip_g.")) { LOG_WARN("loading clip_g from '%s' failed", sd_ctx_params->clip_g_path); } } if (strlen(SAFE_STR(sd_ctx_params->clip_vision_path)) > 0) { LOG_INFO("loading clip_vision from '%s'", sd_ctx_params->clip_vision_path); - if (!model_loader.init_from_file(sd_ctx_params->clip_vision_path, "clip_vision.")) { + if (!load_model("clip_vision", sd_ctx_params->clip_vision_path, "clip_vision.")) { LOG_WARN("loading clip_vision from '%s' failed", sd_ctx_params->clip_vision_path); } } if (strlen(SAFE_STR(sd_ctx_params->t5xxl_path)) > 0) { LOG_INFO("loading t5xxl from '%s'", sd_ctx_params->t5xxl_path); - if (!model_loader.init_from_file(sd_ctx_params->t5xxl_path, "text_encoders.t5xxl.transformer.")) { + if (!load_model("t5xxl", sd_ctx_params->t5xxl_path, "text_encoders.t5xxl.transformer.")) { LOG_WARN("loading t5xxl from '%s' failed", sd_ctx_params->t5xxl_path); } } if (strlen(SAFE_STR(sd_ctx_params->pulid_weights_path)) > 0) { LOG_INFO("loading PuLID weights from '%s'", sd_ctx_params->pulid_weights_path); - if (!model_loader.init_from_file(sd_ctx_params->pulid_weights_path, - "model.diffusion_model.")) { + if (!load_model("pulid_weights", sd_ctx_params->pulid_weights_path, + "model.diffusion_model.")) { LOG_WARN("loading PuLID weights from '%s' failed", sd_ctx_params->pulid_weights_path); } } if (strlen(SAFE_STR(sd_ctx_params->llm_path)) > 0) { LOG_INFO("loading llm from '%s'", sd_ctx_params->llm_path); - if (!model_loader.init_from_file(sd_ctx_params->llm_path, "text_encoders.llm.")) { + if (!load_model("llm", sd_ctx_params->llm_path, "text_encoders.llm.")) { LOG_WARN("loading llm from '%s' failed", sd_ctx_params->llm_path); } } if (strlen(SAFE_STR(sd_ctx_params->llm_vision_path)) > 0) { LOG_INFO("loading llm vision from '%s'", sd_ctx_params->llm_vision_path); - if (!model_loader.init_from_file(sd_ctx_params->llm_vision_path, "text_encoders.llm.visual.")) { + if (!load_model("llm_vision", sd_ctx_params->llm_vision_path, "text_encoders.llm.visual.")) { LOG_WARN("loading llm vision from '%s' failed", sd_ctx_params->llm_vision_path); } } if (strlen(SAFE_STR(sd_ctx_params->vae_path)) > 0) { LOG_INFO("loading vae from '%s'", sd_ctx_params->vae_path); - if (!model_loader.init_from_file(sd_ctx_params->vae_path, "vae.")) { + if (!load_model("vae", sd_ctx_params->vae_path, "vae.")) { LOG_WARN("loading vae from '%s' failed", sd_ctx_params->vae_path); external_vae_is_invalid = true; } @@ -816,7 +770,7 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->taesd_path)) > 0) { LOG_INFO("loading tae from '%s'", sd_ctx_params->taesd_path); - if (!model_loader.init_from_file(sd_ctx_params->taesd_path, "tae.")) { + if (!load_model("taesd", sd_ctx_params->taesd_path, "tae.")) { LOG_WARN("loading tae from '%s' failed", sd_ctx_params->taesd_path); } else { use_tae = true; @@ -825,14 +779,14 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->embeddings_connectors_path)) > 0) { LOG_INFO("loading embeddings connectors from '%s'", sd_ctx_params->embeddings_connectors_path); - if (!model_loader.init_from_file(sd_ctx_params->embeddings_connectors_path)) { + if (!load_model("embeddings_connectors", sd_ctx_params->embeddings_connectors_path, "embeddings_connectors.")) { LOG_WARN("loading embeddings connectors from '%s' failed", sd_ctx_params->embeddings_connectors_path); } } if (strlen(SAFE_STR(sd_ctx_params->audio_vae_path)) > 0) { LOG_INFO("loading audio VAE from '%s'", sd_ctx_params->audio_vae_path); - if (!model_loader.init_from_file(sd_ctx_params->audio_vae_path)) { + if (!load_model("audio_vae", sd_ctx_params->audio_vae_path)) { LOG_WARN("loading audio VAE weights from '%s' failed", sd_ctx_params->audio_vae_path); } else { use_audio_vae = true; @@ -841,15 +795,15 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->audio_encoder_path)) > 0) { LOG_INFO("loading audio encoder (wav2vec2) from '%s'", sd_ctx_params->audio_encoder_path); - if (!model_loader.init_from_file(sd_ctx_params->audio_encoder_path, "wav2vec2.")) { + if (!load_model("audio_encoder", sd_ctx_params->audio_encoder_path, "wav2vec2.")) { LOG_WARN("loading audio encoder weights from '%s' failed", sd_ctx_params->audio_encoder_path); } } if (strlen(SAFE_STR(sd_ctx_params->motion_module_path)) > 0) { LOG_INFO("loading motion module (AnimateDiff) from '%s'", sd_ctx_params->motion_module_path); - if (!model_loader.init_from_file(sd_ctx_params->motion_module_path, - "model.diffusion_model.motion_module.")) { + if (!load_model("motion_module", sd_ctx_params->motion_module_path, + "model.diffusion_model.motion_module.")) { LOG_WARN("loading motion module from '%s' failed", sd_ctx_params->motion_module_path); } else { configuration.animatediff_loaded = true; @@ -863,8 +817,11 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf } } + if (configuration.control_net_file) + configuration.record_model("control_net", sd_ctx_params->control_net_path); + if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0) { - if (!model_loader.init_from_file(sd_ctx_params->ip_adapter_path)) { + if (!load_model("ip_adapter", sd_ctx_params->ip_adapter_path)) { LOG_ERROR("init ip-adapter model loader from file failed: '%s'", sd_ctx_params->ip_adapter_path); return false; } @@ -872,6 +829,8 @@ bool StableDiffusionGGML::init_model_loader(ModelLoader& model_loader, ModelConf if (strlen(SAFE_STR(sd_ctx_params->photo_maker_path)) > 0) { configuration.photomaker_source_available = model_loader.add_file(sd_ctx_params->photo_maker_path, "pmid."); + if (configuration.photomaker_source_available) + configuration.record_model("photo_maker", sd_ctx_params->photo_maker_path); if (!configuration.photomaker_source_available) { LOG_WARN("loading stacked ID embedding from '%s' failed", sd_ctx_params->photo_maker_path); } @@ -1014,7 +973,7 @@ bool StableDiffusionGGML::init(const sd_ctx_params_t* sd_ctx_params) { LOG_ERROR("get sd version from file failed: '%s'", SAFE_STR(sd_ctx_params->model_path)); return false; } else { - LOG_INFO("Version: %s ", model_version_to_str[version]); + LOG_INFO("Version: %s ", sd_version_name(version)); } if (auto_fit_enabled) { @@ -1182,6 +1141,27 @@ bool StableDiffusionGGML::build_control_net_runner() { if (config_->control_net_file == 0) { return true; } + if (sd_version_is_z_image(version)) { + auto z_image = std::dynamic_pointer_cast(diffusion_model); + if (z_image == nullptr || !z_image->control_enabled()) { + LOG_WARN("z-image control net file present but the diffusion model has no control weights; ignoring"); + return true; + } + std::map control_tensors; + z_image->get_control_param_tensors(control_tensors); + if (control_tensors.empty()) { + LOG_ERROR("z-image control net has no param tensors to register"); + return false; + } + ModelManager::ResidencyMode residency_mode = backend_manager.params_backend_is_disk(SDBackendModule::DIFFUSION) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend; + return model_manager->register_param_tensors(ModelComponent::ControlNet, + std::move(control_tensors), + residency_mode, + backend_for(SDBackendModule::DIFFUSION), + params_backend_for(SDBackendModule::DIFFUSION)); + } if (!sd::model_builders::build_control_net_runner(model_build_context(), control_net)) { return false; } @@ -1490,97 +1470,16 @@ bool StableDiffusionGGML::build_denoiser() { prediction_t pred_type = sd_ctx_params->prediction; if (pred_type == PREDICTION_COUNT) { - if (sd_version_is_sd2(version)) { + const auto* definition = find_model_definition(version); + if (!definition) + return false; + const auto& policy = definition->sampling; + pred_type = policy.default_prediction; + default_flow_shift = policy.flow_shift_from_weights ? policy.flow_shift_from_weights(tensor_storage_map) : policy.default_flow_shift; + if (policy.probe_v_prediction) { pred_type = is_using_v_parameterization_for_sd2(sd_version_is_inpaint(version)) ? V_PRED : EPS_PRED; - } else if (sd_version_is_sdxl(version)) { - if (tensor_storage_map.find("edm_vpred.sigma_max") != tensor_storage_map.end()) { - // CosXL models - // TODO: get sigma_min and sigma_max values from file - pred_type = EDM_V_PRED; - } else if (tensor_storage_map.find("v_pred") != tensor_storage_map.end()) { - pred_type = V_PRED; - } else { - pred_type = EPS_PRED; - } - } else if (sd_version_is_sd3(version) || - sd_version_is_wan(version) || - sd_version_is_hunyuan_video(version) || - sd_version_is_lingbot_video(version) || - sd_version_is_minimax_h3(version) || - sd_version_is_qwen_image(version) || - sd_version_is_mage_flow(version) || - version == VERSION_HIDREAM_O1 || - sd_version_is_anima(version) || - sd_version_is_ernie_image(version) || - sd_version_is_z_image(version) || - sd_version_is_z_image_l2p(version) || - version == VERSION_MING_IMAGE || - sd_version_is_llada_image(version) || - sd_version_is_boogu_image(version) || - sd_version_is_pid(version) || - sd_version_is_ideogram4(version) || - version == VERSION_IRIS || - version == VERSION_MOSS_IMAGE) { - pred_type = FLOW_PRED; - if (sd_version_is_wan(version) || version == VERSION_MOSS_IMAGE) { - default_flow_shift = 5.f; - } else if (sd_version_is_hunyuan_video(version)) { - default_flow_shift = 7.f; - } else if (sd_version_is_minimax_h3(version)) { - default_flow_shift = 12.f; - } else if (sd_version_is_ernie_image(version) || version == VERSION_IRIS) { - default_flow_shift = 4.f; - } else if (sd_version_is_pid(version)) { - default_flow_shift = 1.5f; - } else if (sd_version_is_ideogram4(version)) { - default_flow_shift = 1.0f; - } else if (sd_version_is_boogu_image(version)) { - default_flow_shift = 3.16f; - } else if (sd_version_is_mage_flow(version)) { - default_flow_shift = 6.f; - } else if (version == VERSION_MING_IMAGE) { - default_flow_shift = INFINITY; - } else if (sd_version_is_llada_image(version)) { - default_flow_shift = 1.0f; // unused: LLADA_IMAGE_SCHEDULER builds a fixed grid - } else { - default_flow_shift = 3.f; - } - } else if (sd_version_is_flux3(version)) { - pred_type = FLUX_FLOW_PRED; - default_flow_shift = 6.93f; - } else if (sd_version_is_flux(version) || - sd_version_is_flux2(version) || - sd_version_is_longcat(version) || - sd_version_is_lens(version) || - sd_version_is_ltxav(version) || - sd_version_is_krea2(version)) { - pred_type = FLUX_FLOW_PRED; - - default_flow_shift = 1.0f; // TODO: validate - for (const auto& [name, tensor_storage] : tensor_storage_map) { - if (starts_with(name, "model.diffusion_model.guidance_in.in_layer.weight")) { - default_flow_shift = 1.15f; - break; - } - } - if (sd_version_is_longcat(version)) { - default_flow_shift = 3.0f; - } else if (sd_version_is_lens(version)) { - default_flow_shift = 1.83f; - } else if (sd_version_is_ltxav(version)) { - default_flow_shift = 2.37f; - } else if (sd_version_is_krea2(version)) { - default_flow_shift = 1.15f; - } - } else if (sd_version_is_sefi_image(version)) { - pred_type = SEFI_FLOW_PRED; - } else if (sd_version_is_minit2i(version)) { - pred_type = MINIT2I_FLOW_PRED; - } else if (sd_version_is_sensenova_u1(version)) { - pred_type = SENSENOVA_U1_FLOW_PRED; - default_flow_shift = 3.f; - } else { - pred_type = EPS_PRED; + } else if (policy.prediction_from_weights) { + pred_type = policy.prediction_from_weights(tensor_storage_map); } } @@ -1597,27 +1496,28 @@ bool StableDiffusionGGML::build_denoiser() { denoiser = std::make_shared(); break; case FLOW_PRED: { - if (sd_version_is_ltxav(version)) { - LOG_INFO("running in LTXAV FLOW mode"); - denoiser = std::make_shared(); - } else if (sd_version_is_minimax_h3(version)) { - LOG_INFO("running in MiniMax H3 AV FLOW mode"); - denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); - } else if (version == VERSION_MING_IMAGE) { - denoiser = std::make_shared(); - } else { - LOG_INFO("running in FLOW mode"); - denoiser = std::make_shared(); - } + LOG_INFO("running in FLOW mode"); + denoiser = std::make_shared(); break; } case FLUX_FLOW_PRED: { LOG_INFO("running in Flux FLOW mode"); - if (sd_version_is_flux3(version)) { - denoiser = std::make_shared(); - } else { - denoiser = std::make_shared(); - } + denoiser = std::make_shared(); + break; + } + case FLUX3_FLOW_PRED: { + LOG_INFO("running in Flux.3 FLOW mode"); + denoiser = std::make_shared(); + break; + } + case H3_AV_FLOW_PRED: { + LOG_INFO("running in MiniMax H3 AV FLOW mode"); + denoiser = std::make_shared(default_flow_shift, version == VERSION_P_H3_IMAGE ? 1.f : 3.f, get_latent_channel()); + break; + } + case MING_IMAGE_FLOW_PRED: { + LOG_INFO("running in Ming-Image FLOW mode"); + denoiser = std::make_shared(); break; } case SEFI_FLOW_PRED: { @@ -1756,8 +1656,8 @@ std::shared_ptr StableDiffusionGGML::load_lora_model(const ModelManag ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend; auto lora = std::make_shared(lora_log_id(lora_spec), backend_for(module), params_backend_for(module), - model_manager, lora_spec.file_id, version, mode, - backend_manager.params_backend_follows_runtime(module)); + model_manager, lora_spec.file_id, version, mode, + backend_manager.params_backend_follows_runtime(module)); LoraModel::filter_t lora_tensor_filter = module_filter; if (!lora_spec.tensor_name_prefix_filter.empty()) { lora_tensor_filter = [module_filter, prefix = lora_spec.tensor_name_prefix_filter](const std::string& tensor_name) { @@ -2171,11 +2071,11 @@ void StableDiffusionGGML::preview_image(int step, : sd::ops::slice(latents, 2, 0, channels) : latents; if (preview_mode == PREVIEW_PROJ) { - int patch_sz = 1; - const float(*latent_rgb_proj)[3] = nullptr; - float* latent_rgb_bias = nullptr; - const float* latent_alpha_proj = nullptr; - float latent_alpha_bias = 1.f; + int patch_sz = 1; + const float (*latent_rgb_proj)[3] = nullptr; + float* latent_rgb_bias = nullptr; + const float* latent_alpha_proj = nullptr; + float latent_alpha_bias = 1.f; if (channels == 128) { if (sd_version_uses_flux2_vae(version)) { @@ -2457,8 +2357,8 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptris_enabled(); }); - using QwenPrefixInputs = std::tuple*, const sd::Tensor*, - const std::vector>*>; + using QwenPrefixInputs = std::tuple*, const sd::Tensor*, + const std::vector>*>; std::vector qwen_prefix_inputs; RunnerEndOnExit sample_control_runner_end{!control_image.empty() && control_net != nullptr ? control_net.get() : nullptr}; @@ -2760,9 +2660,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::sample(const std::shared_ptr StableDiffusionGGML::sample(const std::shared_ptr(&diffusion_params.extra); if (extra != nullptr) { auto key = std::make_tuple(diffusion_params.context, extra->image_slots, - diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr); + diffusion_params.ref_image_params.pass_to_dit ? diffusion_params.ref_latents : nullptr); auto entry = std::find(qwen_prefix_inputs.begin(), qwen_prefix_inputs.end(), key); extra->prefix_id = static_cast(entry - qwen_prefix_inputs.begin()) + 1; if (entry == qwen_prefix_inputs.end()) { @@ -3055,6 +2959,18 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr::max_digits10); + for (size_t i = 0; i < sigmas.size(); ++i) { + if (i != 0) { + oss << ","; + } + oss << sigmas[i]; + } + LOG_DEBUG("sigmas: %s", oss.str().c_str()); + } + auto x0_opt = sample_k_diffusion(method, denoise, x_t, sigmas, sampler_rng, eta, is_flow_denoiser, extra_sample_args, denoiser); if (x0_opt.empty()) { LOG_ERROR("Diffusion model sampling failed"); @@ -3094,63 +3010,15 @@ int StableDiffusionGGML::get_vae_scale_factor() { } int StableDiffusionGGML::get_diffusion_model_down_factor() { - int down_factor = 8; // unet - if (sd_version_is_dit(version)) { - if (sd_version_is_sensenova_u1(version)) { - down_factor = 32; - } else if (sd_version_is_z_image_l2p(version) || version == VERSION_IRIS) { - down_factor = 16; - } else if (version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE || sd_version_is_wan(version) || sd_version_is_lingbot_video(version) || sd_version_is_minimax_h3(version) || sd_version_is_pixart(version)) { - down_factor = 2; - } else { - down_factor = 1; - } - } - return down_factor; + return find_model_definition(version)->geometry.diffusion_down_factor; } int StableDiffusionGGML::get_latent_channel() { - int latent_channel = 4; - if (sd_version_is_dit(version)) { - if (sd_version_is_ltxav(version)) { - latent_channel = 128; - } else if (sd_version_is_minimax_h3(version)) { - latent_channel = 24; - } else if (version == VERSION_QWEN_IMAGE_2_1) { - latent_channel = 64; - } else if (version == VERSION_WAN2_2_TI2V) { - latent_channel = 48; - } else if (sd_version_is_hunyuan_video(version)) { - latent_channel = 32; - } else if (version == VERSION_HIDREAM_O1) { - latent_channel = 3; - } else if (version == VERSION_CHROMA_RADIANCE) { - latent_channel = 3; - } else if (sd_version_is_minit2i(version)) { - latent_channel = 3; - } else if (sd_version_is_sensenova_u1(version)) { - latent_channel = 3; - } else if (sd_version_is_z_image_l2p(version)) { - latent_channel = 3; - } else if (sd_version_is_pid(version) || version == VERSION_IRIS) { - latent_channel = 3; - } else if (sd_version_is_sefi_image(version)) { - latent_channel = 144; - } else if (sd_version_uses_flux2_vae(version)) { - latent_channel = 128; - } else if (sd_version_is_mage_flow(version)) { - latent_channel = 128; - } else if (sd_version_is_pixart(version)) { - latent_channel = 4; - } else { - latent_channel = 16; - } - } - return latent_channel; + return find_model_definition(version)->geometry.latent_channels; } int StableDiffusionGGML::get_image_channels() const { - return version == VERSION_QWEN_IMAGE_LAYERED || version == VERSION_QWEN_IMAGE_2_1 || version == VERSION_MING_IMAGE ? 4 : 3; + return find_model_definition(version)->geometry.image_channels; } int StableDiffusionGGML::get_image_seq_len(int h, int w) { @@ -3291,7 +3159,10 @@ sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensordiffusion_to_vae_latents(x); + auto image_latents = version == VERSION_P_H3_IMAGE && x.dim() == 5 + ? sd::ops::slice(x, 3, 0, get_latent_channel()) + : x; + auto latents = first_stage_model->diffusion_to_vae_latents(image_latents); auto tiling_params = first_stage_model->resolve_tiling_params(vae_tiling_params); const bool prefer_temporal_tiling = decode_video && latents.dim() == 5 && latents.shape()[2] > 1 && first_stage_model->can_temporal_tile_decode(); @@ -3305,6 +3176,9 @@ sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensordecode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y); const bool temporal_before = tiling_params.temporal_tiling; + if (version == VERSION_P_H3_IMAGE && !decoded.empty() && decoded.dim() == 5) { + decoded = decoded.squeeze(2); + } if (!decoded.empty() || !sd::backend_fit::prepare_vae_decode_retry_tiling(tiling_params, prefer_temporal_tiling, first_stage_model->last_compute_status(), @@ -3359,38 +3233,15 @@ bool StableDiffusionGGML::is_flow_denoiser() { } std::string StableDiffusionGGML::get_default_ref_image_preset(SDVersion version) const { - if (sd_version_is_longcat(version)) { - return "longcat"; - } else if (sd_version_is_flux(version)) { - return "flux_kontext"; - } else if (sd_version_is_flux3(version)) { - return "flux2"; - } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { - return "flux2"; - } else if (version == VERSION_QWEN_IMAGE_LAYERED) { - return "qwen_layered"; - } else if (sd_version_is_qwen_image(version)) { - return "qwen"; - } else if (sd_version_is_mage_flow(version)) { - return "mage_flow"; - } else if (sd_version_is_z_image(version) || sd_version_is_boogu_image(version)) { - return "z_image_omni"; - } else if (sd_version_is_llada_image(version)) { - return "llada_image"; - } else if (sd_version_is_krea2(version)) { - // have to make a choice between "krea2_edit" mode (for lbouaraba/krea2edit) - // and "krea2_ostris_edit" (for krea2 ostris edit) - // since krea2 ostris edit support predates, it should probably be default - return "krea2_ostris_edit"; - } else if (sd_version_is_anima(version)) { - return "cosmos_reference"; - } - return "default"; + return find_model_definition(version)->ref_image_preset; } RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_image_args) const { RefImageParams params; std::string preset_name = get_default_ref_image_preset(version); + // Keys whose value was rejected. They are already reported, so per-model "no effect" + // warnings must not repeat them. + std::set rejected_keys; for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { if (key == "preset") { @@ -3403,6 +3254,7 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima valid_list += (valid_list.empty() ? "" : ", ") + name; } LOG_WARN("ignoring invalid reference image preset '%s'. Valid options: [%s]", value.c_str(), valid_list.c_str()); + rejected_keys.insert(key); } break; } @@ -3411,15 +3263,20 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima LOG_INFO("Using '%s' preset for reference images", preset_name.c_str()); params = REF_IMAGE_PRESETS.at(preset_name); } + if (version == VERSION_QWEN_IMAGE_2_1) { + params.vlm_use_vae_input = true; + } for (const auto& [key, value] : parse_key_value_args(ref_image_args, "reference image args")) { if (key == "pass_to_vlm") { if (!parse_strict_bool(value, params.pass_to_vlm)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "pass_to_dit") { if (!parse_strict_bool(value, params.pass_to_dit)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "ref_index_mode") { if (value == "fixed") { @@ -3430,18 +3287,42 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima params.ref_index_mode = Rope::RefIndexMode::DECREASE; } else { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "force_ref_timestep_zero") { if (!parse_strict_bool(value, params.force_ref_timestep_zero)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "resize_before_vae") { if (!parse_strict_bool(value, params.resize_before_vae)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); + } + } else if (key == "resize_vae_to_target") { + if (!parse_strict_bool(value, params.resize_vae_to_target)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "resize_single_ref_to_target") { + if (!parse_strict_bool(value, params.resize_single_ref_to_target)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vae_size_buckets") { + if (!parse_strict_bool(value, params.vae_size_buckets)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vae_center_crop") { + if (!parse_strict_bool(value, params.vae_center_crop)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + } + } else if (key == "vlm_use_vae_input") { + if (!parse_strict_bool(value, params.vlm_use_vae_input)) { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); } } else if (key == "vae_input_max_pixels") { if (!parse_strict_int(value, params.vae_input_max_pixels)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "vlm_resize_mode") { if (value == "longest_side") { @@ -3452,14 +3333,31 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima params.vlm_resize_mode = RefImageResizeMode::NONE; } else { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "vlm_max_size") { if (!parse_strict_int(value, params.vlm_max_size)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } } else if (key == "vlm_min_size") { if (!parse_strict_int(value, params.vlm_min_size)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); + } + } else if (key == "ref_image_size") { + bool valid = true; + if (value == "match") { + params.ref_image_size = RefImageSizeMode::MATCH; + } else if (value == "max") { + params.ref_image_size = RefImageSizeMode::MAX; + } else { + LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); + valid = false; + } + if (valid && !sd_version_is_minimax_h3(version)) { + LOG_WARN("ref_image_size is only supported by MiniMax-H3 for now"); } } else if (key != "preset" && key != "vlm_size") { LOG_WARN("ignoring unknown reference image arg '%s'", key.c_str()); @@ -3470,6 +3368,7 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima int vlm_size; if (!parse_strict_int(value, vlm_size)) { LOG_WARN("ignoring invalid reference image arg '%s=%s'", key.c_str(), value.c_str()); + rejected_keys.insert(key); } else { LOG_INFO("vlm_size override: setting both min and max size to %ld", (long)vlm_size); params.vlm_min_size = vlm_size; @@ -3481,6 +3380,29 @@ RefImageParams StableDiffusionGGML::resolve_ref_image_params(const char* ref_ima if (params.force_ref_timestep_zero && !sd_version_is_krea2(version)) { LOG_WARN("force_ref_timestep_zero is only supported by Krea2 architecture for now"); } + if (version == VERSION_MINIMAX_H3) { + // H3 consumes references as VAE latents in its own packed sequence rather than + // through the generic conditioner/DiT reference path, so ref_image_size is the + // only arg it acts on. Everything else is reported as ineffective rather than + // enumerated, so newly added RefImageParams members cannot silently do nothing. + static const std::set h3_effective_keys = {"ref_image_size"}; + // Keys that already have their own, more specific warning above. + static const std::set h3_separately_reported_keys = { + "resize_before_vae", + "force_ref_timestep_zero", + }; + for (const auto& arg : parse_key_value_args(ref_image_args, "reference image args")) { + if (h3_effective_keys.count(arg.first) == 0 && + h3_separately_reported_keys.count(arg.first) == 0 && + rejected_keys.count(arg.first) == 0) { + LOG_WARN("reference image arg '%s' has no effect for MiniMax-H3", arg.first.c_str()); + } + } + + if (!params.resize_before_vae) { + LOG_WARN("MiniMax-H3 requires 32-pixel-aligned reference images and cannot skip resizing before VAE encoding; ignoring resize_before_vae=false"); + } + } return params; } diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index c5509d9a9e..e067781b45 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -34,10 +34,9 @@ namespace Wav2Vec2 { class Wav2Vec2ModelRunner; } -extern const char* model_version_to_str[]; - static inline bool sd_version_supports_ref_latent_img_cfg(SDVersion version) { return version == VERSION_FLUX || + version == VERSION_CHROMA || sd_version_is_flux2(version) || (sd_version_is_qwen_image(version) && version != VERSION_QWEN_IMAGE_2_1) || sd_version_is_mage_flow(version) || @@ -125,6 +124,7 @@ class StableDiffusionGGML { sd_ctx_params_t params{}; std::list strings; std::vector embeddings; + std::map models; ModelLoader::FileId control_net_file = 0; bool use_tae = false; bool use_audio_vae = false; @@ -166,13 +166,24 @@ class StableDiffusionGGML { use_audio_vae = other.use_audio_vae; photomaker_source_available = other.photomaker_source_available; animatediff_loaded = other.animatediff_loaded; + models = other.models; } ModelConfig& operator=(const ModelConfig&) = delete; + void record_model(const std::string& name, const std::string& path) { + if (path.empty()) { + models.erase(name); + return; + } + const auto separator = path.find_last_of("/\\"); + models[name] = separator == std::string::npos ? path : path.substr(separator + 1); + } + void set_control_net(ModelLoader::FileId id, const std::string& path) { control_net_file = id; strings.push_back(path); params.control_net_path = strings.back().c_str(); + record_model("control_net", path); } }; diff --git a/src/pipeline/generation.h b/src/pipeline/generation.h index 34aa000aa6..0bc3913150 100644 --- a/src/pipeline/generation.h +++ b/src/pipeline/generation.h @@ -2,6 +2,7 @@ #define __SD_PIPELINE_GENERATION_H__ #include "conditioning/conditioner.hpp" +#include "json.hpp" #include "stable-diffusion.h" class StableDiffusionGGML; @@ -54,14 +55,16 @@ namespace sd::pipeline { bool generate_image(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, sd_image_t** images_out, - int* num_images_out); + int* num_images_out, + std::vector* infos_out = nullptr); bool generate_video(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, sd_image_t** frames_out, int* num_frames_out, sd_audio_t** audio_out, - int* fps_out); + int* fps_out, + nlohmann::json* info_out = nullptr); sd::Tensor upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, const char* model_path, diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 7de85b3595..fc3adb9639 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -1,4 +1,5 @@ #include "generation.h" +#include "generation_info.h" #include #include @@ -119,6 +120,60 @@ namespace sd::pipeline { channels); } + static sd::Tensor encode_z_image_control_latent(StableDiffusionGGML* sd, const sd::Tensor& control_rgb) { + sd::Tensor vae_output = sd->first_stage_model->encode(sd->n_threads, + control_rgb, + sd->vae_tiling_params, + sd->circular_x, + sd->circular_y); + if (vae_output.empty()) { + LOG_ERROR("failed to encode control image"); + return {}; + } + // TAE encode already returns the latent; KL/Flux VAE returns a mean/logvar pair. + sd::Tensor ctrl_latent = sd->first_stage_model->get_desc() == "taesd" + ? vae_output + : sd::ops::chunk(vae_output, 2, 2)[0]; + ctrl_latent = sd->first_stage_model->vae_to_diffusion_latents(ctrl_latent); + if (ctrl_latent.empty() || ctrl_latent.dim() != 4 || ctrl_latent.shape()[2] > 33) { + LOG_ERROR("failed to encode control image"); + return {}; + } + std::vector padded_shape = ctrl_latent.shape(); + padded_shape[2] = 33; + sd::Tensor padded(padded_shape); + sd::ops::slice_assign(&padded, 2, 0, ctrl_latent.shape()[2], ctrl_latent); + return padded; + } + + static sd::Tensor resize_ph3_reference(const sd::Tensor& image, int64_t pixels) { + double scale = std::min(1.0, std::sqrt(static_cast(pixels) / (image.shape()[0] * image.shape()[1]))); + int64_t width = std::max(32, static_cast(std::nearbyint(image.shape()[0] * scale / 32)) * 32); + int64_t height = std::max(32, static_cast(std::nearbyint(image.shape()[1] * scale / 32)) * 32); + if (width == image.shape()[0] && height == image.shape()[1]) { + return image; + } + return sd::ops::interpolate(image, {width, height, image.shape()[2], image.shape()[3]}, sd::ops::InterpolateMode::Lanczos); + } + + static sd::Tensor pack_ph3_image_latent(const sd::Tensor& image, + const sd::Tensor& previous = {}) { + auto video = image.dim() == 4 ? image.unsqueeze(2) : image; + GGML_ASSERT(video.dim() == 5 && video.shape()[2] == 1 && video.shape()[3] == 24 && video.shape()[4] == 1); + const int64_t spatial = video.shape()[0] * video.shape()[1]; + // The image DiT still attends to two stereo audio latents (32 channels). + auto shape = video.shape(); + shape[3] += (128 + spatial - 1) / spatial; + auto packed = sd::zeros(shape); + std::copy_n(video.data(), video.numel(), packed.data()); + if (!previous.empty()) { + const int64_t offset = previous.shape()[0] * previous.shape()[1] * 24; + GGML_ASSERT(previous.numel() >= offset + 128); + std::copy_n(previous.data() + offset, 128, packed.data() + video.numel()); + } + return packed; + } + static std::optional prepare_image_generation_latents(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, GenerationRequest* request, @@ -252,10 +307,22 @@ namespace sd::pipeline { } if (!control_image_tensor.empty()) { - control_latent = sd->encode_first_stage(control_image_tensor); - if (control_latent.empty()) { - LOG_ERROR("failed to encode control image"); - return std::nullopt; + const bool z_image_control = sd_version_is_z_image(sd->version) && + sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0; + if (z_image_control) { + control_image_tensor = encode_z_image_control_latent(sd, control_image_tensor); + if (control_image_tensor.empty()) { + return std::nullopt; + } + } else if (sd_version_is_z_image(sd->version)) { + LOG_WARN("control image provided but the z-image model has no control net; ignoring"); + control_image_tensor = {}; + } else { + control_latent = sd->encode_first_stage(control_image_tensor); + if (control_latent.empty()) { + LOG_ERROR("failed to encode control image"); + return std::nullopt; + } } } @@ -276,36 +343,51 @@ namespace sd::pipeline { LOG_INFO("EDIT mode"); } + bool skip_vae_encoding = (sd->version == VERSION_HIDREAM_O1) || !( + ref_image_params.pass_to_dit || ref_image_params.vlm_use_vae_input || + sd_version_is_unet_edit(sd->version) || sd_version_is_pid(sd->version)); + std::vector> ref_latents; for (size_t i = 0; i < ref_images.size(); i++) { - if (sd->version == VERSION_HIDREAM_O1) { + if (skip_vae_encoding) { continue; } sd::Tensor ref_latent; - if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd->version)) { + if (sd->version == VERSION_P_H3_IMAGE) { + auto resized = resize_ph3_reference(ref_images[i], int64_t(request->width) * request->height); + ref_latent = sd->encode_first_stage(resized); + } else if (ref_image_params.resize_before_vae && !sd_version_is_pid(sd->version)) { LOG_VERBOSE("auto resize ref images"); double vae_width; double vae_height; - if (ref_image_params.resize_vae_to_target) { + int factor = sd_version_is_qwen_image(sd->version) ? 32 : 16; + if (ref_image_params.resize_vae_to_target || (ref_image_params.resize_single_ref_to_target && ref_images.size() == 1)) { vae_width = request->width; vae_height = request->height; } else { int default_pixels = sd->version == VERSION_QWEN_IMAGE_2_1 ? request->width * request->height : 1024 * 1024; int target_pixels = ref_image_params.vae_input_max_pixels > 0 ? ref_image_params.vae_input_max_pixels : default_pixels; - int vae_image_size = std::min(target_pixels, request->width * request->height); - vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); - vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + if (ref_image_params.vae_size_buckets) { + auto size = sd::select_image_size_bucket(static_cast(ref_images[i].shape()[0]), static_cast(ref_images[i].shape()[1]), target_pixels, factor, request->width, request->height); + vae_width = size.first; + vae_height = size.second; + } else { + int vae_image_size = std::min(target_pixels, request->width * request->height); + vae_width = sqrt(vae_image_size * ref_images[i].shape()[0] / ref_images[i].shape()[1]); + vae_height = vae_width * ref_images[i].shape()[1] / ref_images[i].shape()[0]; + } } - int factor = sd_version_is_qwen_image(sd->version) ? 32 : 16; vae_height = round(vae_height / factor) * factor; vae_width = round(vae_width / factor) * factor; - auto resized_ref_img = sd::ops::interpolate(ref_images[i], - {static_cast(vae_width), - static_cast(vae_height), - ref_images[i].shape()[2], - ref_images[i].shape()[3]}); + auto resized_ref_img = ref_image_params.vae_center_crop + ? sd::resize_and_center_crop(ref_images[i], static_cast(vae_width), static_cast(vae_height)) + : sd::ops::interpolate(ref_images[i], + {static_cast(vae_width), + static_cast(vae_height), + ref_images[i].shape()[2], + ref_images[i].shape()[3]}); LOG_VERBOSE("resize vae ref image %d from %" PRId64 "x%" PRId64 " to %" PRId64 "x%" PRId64, static_cast(i), @@ -315,7 +397,7 @@ namespace sd::pipeline { resized_ref_img.shape()[0]); ref_latent = sd->encode_first_stage(resized_ref_img); - if (sd->version == VERSION_QWEN_IMAGE_2_1) { + if (ref_image_params.vlm_use_vae_input) { ref_images[i] = std::move(resized_ref_img); } } else { @@ -401,6 +483,16 @@ namespace sd::pipeline { } ImageGenerationLatents latents; + if (sd->version == VERSION_P_H3_IMAGE) { + init_latent = pack_ph3_image_latent(init_latent); + latent_mask = latent_mask.unsqueeze(2); + for (size_t i = 0; i < ref_latents.size(); ++i) { + latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE, static_cast(i), -1}); + } + for (auto& image : ref_images) { + image = resize_ph3_reference(image, 1000000); + } + } latents.init_latent = std::move(init_latent); latents.concat_latent = std::move(concat_latent); latents.img_uncond_concat_latent = std::move(img_uncond_concat_latent); @@ -453,6 +545,9 @@ namespace sd::pipeline { if (cond.c_concat.empty() && ref_image_params.pass_to_dit) { cond.c_concat = latents->concat_latent; // TODO: optimize } + if (sd->version == VERSION_P_H3_IMAGE && ref_image_params.pass_to_dit) { + cond.c_reference_blocks = latents->minimax_reference_blocks; + } bool use_ref_latent_img_cfg = request->use_img_uncond && !latents->ref_images.empty() && @@ -484,8 +579,8 @@ namespace sd::pipeline { } condition_params.text = request->negative_prompt; condition_params.zero_out_masked = zero_out_masked; - if (sd_version_is_llada_image(sd->version)) { - // LLaDA-Image CFG keeps the source latent but drops its SigVQ features. + if (sd_version_is_llada_image(sd->version) || sd->version == VERSION_FIBO2) { + // Keep source latents in CFG, but omit source images from the text encoder. condition_params.ref_images = nullptr; } uncond = sd->get_learned_condition(condition_params); @@ -497,6 +592,9 @@ namespace sd::pipeline { if (uncond.c_concat.empty() && ref_image_params.pass_to_dit) { uncond.c_concat = latents->concat_latent; // TODO: optimize } + if (sd->version == VERSION_P_H3_IMAGE && ref_image_params.pass_to_dit) { + uncond.c_reference_blocks = latents->minimax_reference_blocks; + } } SDCondition img_uncond; @@ -624,7 +722,7 @@ namespace sd::pipeline { int64_t t4 = ggml_time_ms(); LOG_INFO("decode_first_stage completed, taking %.2fs", (t4 - t0) * 1.0f / 1000); - if (decoded_images.empty()) { + if (decoded_images.empty() || cancelled) { LOG_ERROR(cancelled ? "cancelled before any latent images were decoded" : "no decoded images"); return nullptr; } @@ -785,7 +883,10 @@ namespace sd::pipeline { bool generate_image(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, sd_image_t** images_out, - int* num_images_out) { + int* num_images_out, + std::vector* infos_out) { + if (infos_out) + infos_out->clear(); if (images_out != nullptr) { *images_out = nullptr; } @@ -798,7 +899,7 @@ namespace sd::pipeline { if (!sd_version_supports_image_generation(sd->version)) { LOG_ERROR("%s cannot be run with generate_image(); use generate_video() or --mode vid_gen in the CLI", - model_version_to_str[sd->version]); + sd_version_name(sd->version)); return false; } @@ -807,6 +908,7 @@ namespace sd::pipeline { int64_t t0 = ggml_time_ms(); sd->vae_tiling_params = sd_img_gen_params->vae_tiling_params; GenerationRequest request(sd, sd_img_gen_params); + const auto info = sd::generation_info(request, *sd_img_gen_params); sd::ImagePreprocessor preprocessing(sd_img_gen_params->image_preprocess.rules); sd_img_gen_params_t processed_params = *sd_img_gen_params; if (!preprocessing.prepare_inputs(processed_params, request.width, request.height)) @@ -847,7 +949,6 @@ namespace sd::pipeline { return false; } ImageGenerationEmbeds embeds = std::move(*embeds_opt); - std::vector> final_latents; int64_t denoise_start = ggml_time_ms(); for (int b = 0; b < request.batch_count; b++) { @@ -864,7 +965,7 @@ namespace sd::pipeline { } int64_t sampling_start = ggml_time_ms(); - int64_t cur_seed = request.seed + b; + int64_t cur_seed = request.seed_for_batch(b); LOG_INFO("generating image: %i/%i - seed %" PRId64, b + 1, request.batch_count, cur_seed); sd->rng->manual_seed(cur_seed); @@ -933,10 +1034,10 @@ namespace sd::pipeline { } LOG_INFO("hires fix: loading model upscaler from '%s'", request.hires.model_path); hires_upscaler = std::make_unique(sd->n_threads, - false, - request.hires.upscale_tile_size, - sd->backend_spec, - sd->params_backend_spec); + false, + request.hires.upscale_tile_size, + sd->backend_spec, + sd->params_backend_spec); const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); hires_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (!hires_upscaler->load_from_file(request.hires.model_path, @@ -950,7 +1051,7 @@ namespace sd::pipeline { std::vector hires_sigma_sched = make_hires_sigma_schedule(sd, request.hires, - sd_img_gen_params->sample_params, + request.sample_params, plan.sample_method, plan.sample_steps, sd->get_image_seq_len(request.hires.target_height, request.hires.target_width), @@ -963,22 +1064,51 @@ namespace sd::pipeline { std::vector> hires_final_latents; int64_t hires_denoise_start = ggml_time_ms(); + sd::Tensor hires_control_image; + if (!latents.control_image.empty() && + sd_version_is_z_image(sd->version) && + sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0) { + const int64_t hires_latent_w = request.hires.target_width / request.vae_scale_factor; + const int64_t hires_latent_h = request.hires.target_height / request.vae_scale_factor; + if (latents.control_image.shape()[0] != hires_latent_w || + latents.control_image.shape()[1] != hires_latent_h) { + LOG_INFO("hires fix: re-encoding control latent for %dx%d", request.hires.target_width, request.hires.target_height); + hires_control_image = encode_z_image_control_latent(sd, + sd_image_to_tensor(sd_img_gen_params->control_image, + request.hires.target_width, + request.hires.target_height)); + if (hires_control_image.empty()) { + std::vector ctrl_shape = latents.control_image.shape(); + ctrl_shape[0] = hires_latent_w; + ctrl_shape[1] = hires_latent_h; + hires_control_image = sd::ops::interpolate(latents.control_image, ctrl_shape, sd::ops::InterpolateMode::Bicubic); + } + } else { + hires_control_image = latents.control_image; + } + } for (int b = 0; b < (int)final_latents.size(); b++) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { LOG_ERROR("cancelling generation during hires fix"); return false; } - int64_t cur_seed = request.seed + b; + int64_t cur_seed = request.seed_for_batch(b); sd->rng->manual_seed(cur_seed); sd->sampler_rng->manual_seed(cur_seed); + auto hires_source = sd->version == VERSION_P_H3_IMAGE + ? sd::ops::slice(final_latents[b], 3, 0, 24) + : final_latents[b]; sd::Tensor upscaled = upscale_hires_latent(sd, - final_latents[b], + hires_source, request, hires_upscaler.get()); if (upscaled.empty()) { return false; } + if (sd->version == VERSION_P_H3_IMAGE) { + upscaled = pack_ph3_image_latent(upscaled, final_latents[b]); + } sd::Tensor noise = sd::randn_like(upscaled, sd->rng); @@ -1000,7 +1130,7 @@ namespace sd::pipeline { embeds.cond, embeds.uncond, embeds.img_uncond, - latents.control_image, + !hires_control_image.empty() ? hires_control_image : latents.control_image, request.control_strength, request.guidance, plan.eta, @@ -1046,6 +1176,19 @@ namespace sd::pipeline { return false; } + if (infos_out) { + try { + infos_out->assign(num_images, info); + // Layers and frames decoded from the same latent share its seed. + const int images_per_latent = num_images / static_cast(final_latents.size()); + for (int i = 0; i < num_images; ++i) { + (*infos_out)[i]["seed"] = std::to_string(request.seed_for_batch(i / images_per_latent)); + } + } catch (...) { + free_sd_images(result, num_images); + throw; + } + } sd->lora_stat(); int64_t t1 = ggml_time_ms(); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 4b435170f7..e78a0d666e 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -10,37 +10,8 @@ #include "extensions/generation_extension.h" #include "model/adapter/ip_adapter.hpp" #include "model/audio/wav2vec2.hpp" -#include "model/diffusion/anima.hpp" -#include "model/diffusion/boogu.hpp" #include "model/diffusion/control.hpp" -#include "model/diffusion/ernie_image.hpp" -#include "model/diffusion/flux.hpp" -#include "model/diffusion/flux3.hpp" -#include "model/diffusion/hidream_o1.hpp" -#include "model/diffusion/hunyuan.hpp" -#include "model/diffusion/ideogram4.hpp" -#include "model/diffusion/iris.hpp" -#include "model/diffusion/krea2.hpp" -#include "model/diffusion/lens.hpp" -#include "model/diffusion/lingbot_video.hpp" -#include "model/diffusion/llada_image.hpp" -#include "model/diffusion/ltxv.hpp" -#include "model/diffusion/mage_flow.hpp" -#include "model/diffusion/ming_image.hpp" -#include "model/diffusion/minimax_h3.hpp" -#include "model/diffusion/minit2i.hpp" -#include "model/diffusion/mmdit.hpp" #include "model/diffusion/model.hpp" -#include "model/diffusion/moss_image.hpp" -#include "model/diffusion/pid.hpp" -#include "model/diffusion/pixart.hpp" -#include "model/diffusion/qwen_image.hpp" -#include "model/diffusion/qwen_image_2_1.hpp" -#include "model/diffusion/sensenova_u1.h" -#include "model/diffusion/unet.hpp" -#include "model/diffusion/wan.hpp" -#include "model/diffusion/z_image.hpp" -#include "model/diffusion/z_image_l2p.hpp" #include "model/vae/auto_encoder_kl.hpp" #include "model/vae/hunyuan_vae.hpp" #include "model/vae/ltx_audio_vae.hpp" @@ -51,10 +22,11 @@ #include "model/vae/tae.hpp" #include "model/vae/vae.hpp" #include "model/vae/wan_vae.hpp" +#include "model_definition.h" namespace sd::model_builders { - static bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) { + bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module) { if (backends.runtime_backend(module) == nullptr) { LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module)); return false; @@ -94,402 +66,10 @@ namespace sd::model_builders { return false; } - if (sd_version_is_sd3(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_pid(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model.net", - weight_manager); - } else if (sd_version_is_ideogram4(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_krea2(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_flux(version)) { - bool is_chroma = false; - for (auto pair : tensor_storage_map) { - if (pair.first.find("distilled_guidance_layer.in_proj.weight") != std::string::npos) { - is_chroma = true; - break; - } - } - if (is_chroma) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - false, - 1, - false, - weight_manager, - sd_ctx_params->model_args); - } else if (version == VERSION_OVIS_IMAGE) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - } else { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - } - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_flux3(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, version, "", false, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (sd_version_is_flux2(version) || sd_version_is_sefi_image(version)) { - bool is_chroma = false; - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_ltxav(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - "text_encoders.llm", - "text_embedding_projection", - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_minimax_h3(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_hunyuan_video(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - } else if (sd_version_is_wan(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - true, - 0, - true, - weight_manager); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - if (strlen(SAFE_STR(sd_ctx_params->high_noise_diffusion_model_path)) > 0) { - result.high_noise_diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.high_noise_diffusion_model", - version, - weight_manager); - } - if (result.diffusion->get_desc() == "Wan2.1-I2V-14B" || - result.diffusion->get_desc() == "Wan2.1-FLF2V-14B" || - result.diffusion->get_desc() == "Wan2.1-I2V-1.3B") { - if (!ensure_backend_pair(ctx.backends, SDBackendModule::CLIP_VISION)) { - return false; - } - result.clip_vision = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::CLIP_VISION), - tensor_storage_map, - weight_manager); - } - if (version == VERSION_WAN2_2_S2V && - tensor_storage_map.count("wav2vec2.encoder.layer_norm.bias") > 0) { - if (!ensure_backend_pair(ctx.backends, SDBackendModule::AUDIO_ENCODER)) { - return false; - } - result.audio_encoder = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::AUDIO_ENCODER), - tensor_storage_map, - "wav2vec2.", - weight_manager); - } - } else if (sd_version_is_lingbot_video(version)) { - bool enable_vision = false; - for (const auto& [name, _] : tensor_storage_map) { - if (starts_with(name, "text_encoders.llm.visual.")) { - enable_vision = true; - break; - } - } - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - enable_vision, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_qwen_image(version)) { - bool enable_vision = version != VERSION_QWEN_IMAGE_LAYERED; - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - enable_vision, - weight_manager, - tokenizers); - if (version == VERSION_QWEN_IMAGE_2_1) { - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager, - sd_ctx_params->model_args); - } else { - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } - } else if (version == VERSION_PIXART) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - true, - 0, - false, - weight_manager, - sd_ctx_params->model_args); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager, - sd_ctx_params->model_args); - } else if (sd_version_is_mage_flow(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_longcat(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager, - sd_ctx_params->model_args); - } else if (version == VERSION_HIDREAM_O1) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model", - weight_manager); - } else if (sd_version_is_minit2i(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model.model.net", - weight_manager); - } else if (sd_version_is_sensenova_u1(version)) { - result.conditioner = std::make_shared(); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "", - weight_manager); - } else if (sd_version_is_anima(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (version == VERSION_MING_IMAGE) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (sd_version_is_z_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - } else if (sd_version_is_z_image_l2p(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_llada_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - "text_encoders.llm", - "queryformer", - "text_projection", - "sigvq", - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_boogu_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - true, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); - } else if (sd_version_is_ernie_image(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (sd_version_is_lens(version)) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - version, - "", - false, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - weight_manager); - } else if (version == VERSION_IRIS) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, version, "", false, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else if (version == VERSION_MOSS_IMAGE) { - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, version, "", false, weight_manager, tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, "model.diffusion_model", weight_manager); - } else { // SD1.x SD2.x SDXL - std::map embbeding_map; - for (uint32_t i = 0; i < sd_ctx_params->embedding_count; i++) { - embbeding_map.emplace(SAFE_STR(sd_ctx_params->embeddings[i].name), SAFE_STR(sd_ctx_params->embeddings[i].path)); - } - result.conditioner = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::TE), - tensor_storage_map, - embbeding_map, - version, - weight_manager, - tokenizers); - result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), - tensor_storage_map, - "model.diffusion_model", - version, - weight_manager); + const auto* definition = find_model_definition(version); + if (!definition || !definition->build || !definition->build(ctx, tokenizers, result)) { + LOG_ERROR("failed to build model runners for %s", sd_version_name(version)); + return false; } if (strlen(SAFE_STR(sd_ctx_params->ip_adapter_path)) > 0 && result.clip_vision == nullptr) { @@ -530,7 +110,7 @@ namespace sd::model_builders { return false; } - bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) { + bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners) try { const auto* sd_ctx_params = &ctx.params; const auto& tensor_storage_map = ctx.tensor_storage_map; const auto version = ctx.version; @@ -541,7 +121,7 @@ namespace sd::model_builders { } auto create_tae = [&](bool decode_only) -> std::shared_ptr { - if ((sd_version_uses_wan_vae(version) && version != VERSION_QWEN_IMAGE_2_1) || sd_version_is_hunyuan_video(version) || sd_version_is_ltxav(version) || sd_version_is_minimax_h3(version)) { + if (find_model_traits(version)->video_tae) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "decoder", @@ -566,36 +146,43 @@ namespace sd::model_builders { vae_format = SD_VAE_FORMAT_AUTO; } SDVersion vae_version = version; - if (sd_version_is_pid(version) && vae_format != SD_VAE_FORMAT_AUTO) { + if (find_model_traits(version)->vae_format_override && vae_format != SD_VAE_FORMAT_AUTO) { vae_version = sd_vae_format_to_version(vae_format, vae_version); } auto create_vae = [&]() -> std::shared_ptr { - if (sd_version_is_ltxav(version)) { + if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::LTXAV) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", false, version, weight_manager); - } else if (sd_version_is_minimax_h3(version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MINIMAX_H3) { + if (version == VERSION_P_H3_IMAGE) { + // The video VAE also loads successfully, but uses different temporal decoding and posterior sampling. + auto proj = tensor_storage_map.find("first_stage_model.decoder.proj_out.weight"); + if (proj == tensor_storage_map.end() || proj->second.ne[1] != 768) { + throw std::runtime_error("P-H3-Image requires its single-frame VAE; pass --vae p-h3-image_vae_bf16.safetensors"); + } + } return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", weight_manager); - } else if (sd_version_is_mage_flow(vae_version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MAGE) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", weight_manager); - } else if (sd_version_uses_hunyuan_video_vae(vae_version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::HUNYUAN) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", false, vae_version, weight_manager); - } else if (sd_version_uses_wan_vae(vae_version)) { + } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::WAN) { return std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "first_stage_model", @@ -603,49 +190,22 @@ namespace sd::model_builders { vae_version, weight_manager); } else { - auto model = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), - tensor_storage_map, - "first_stage_model", - false, - false, - vae_version, - weight_manager); - if (sd_version_is_pixart(version)) { - // Alpha-512 and Sigma share tensor layouts; Alpha-512 needs an explicit scale override. - if (tensor_storage_map.count("model.diffusion_model.csize_embedder.mlp.0.weight") != 0) { - model->scale_factor = 0.18215f; - } - for (const auto& [key, value] : parse_key_value_args(sd_ctx_params->model_args, "model arg")) { - if (key == "pixart_vae_scale_factor") { - float parsed = 0.f; - if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed > 0.f) { - model->scale_factor = parsed; - } else { - LOG_WARN("ignoring invalid PixArt model arg '%s=%s'", key.c_str(), value.c_str()); - } - } - } - LOG_VERBOSE("pixart: VAE scale factor = %.5f", model->scale_factor); - } - if (sd_version_is_sdxl(version) && - (strlen(SAFE_STR(sd_ctx_params->vae_path)) == 0 || sd_ctx_params->force_sdxl_vae_conv_scale || options.external_vae_is_invalid)) { - float vae_conv_2d_scale = 1.f / 32.f; - LOG_WARN( - "No valid VAE specified with --vae or --force-sdxl-vae-conv-scale flag set, " - "using Conv2D scale %.3f", - vae_conv_2d_scale); - model->set_conv2d_scale(vae_conv_2d_scale); + auto model = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), + tensor_storage_map, + "first_stage_model", + false, + false, + vae_version, + weight_manager); + const auto* definition = find_model_definition(version); + if (definition && definition->configure_vae) { + definition->configure_vae(ctx, options, *model); } return model; } }; - if (version == VERSION_CHROMA_RADIANCE || - version == VERSION_HIDREAM_O1 || - sd_version_is_minit2i(version) || - sd_version_is_sensenova_u1(version) || - sd_version_is_z_image_l2p(version) || - version == VERSION_IRIS) { + if (find_model_traits(version)->fake_vae) { LOG_INFO("using FakeVAE"); result.vae = std::make_shared(version, ctx.backends.runtime_backend(SDBackendModule::VAE), @@ -663,7 +223,7 @@ namespace sd::model_builders { } if (options.use_audio_vae) { - if (sd_version_is_minimax_h3(version)) { + if (find_model_traits(version)->vae_kind == ModelVAEKind::MINIMAX_H3) { result.audio = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::VAE), tensor_storage_map, "", @@ -687,6 +247,9 @@ namespace sd::model_builders { } runners = std::move(result); return true; + } catch (const std::exception& error) { + LOG_ERROR("failed to build VAE runners: %s", error.what()); + return false; } bool build_control_net_runner(const Context& ctx, std::shared_ptr& runner) { @@ -694,6 +257,9 @@ namespace sd::model_builders { const auto& tensor_storage_map = ctx.tensor_storage_map; const auto version = ctx.version; const auto& weight_manager = ctx.weight_manager; + if (sd_version_is_z_image(version)) { + return true; + } if (!ensure_backend_pair(ctx.backends, SDBackendModule::CONTROL_NET)) { return false; } diff --git a/src/pipeline/model_builders.h b/src/pipeline/model_builders.h index 5edc75680f..bb51ccc14a 100644 --- a/src/pipeline/model_builders.h +++ b/src/pipeline/model_builders.h @@ -4,6 +4,7 @@ #include #include +#include "core/ggml_extend_backend.h" #include "model.h" #include "stable-diffusion.h" @@ -56,6 +57,8 @@ namespace sd::model_builders { std::shared_ptr audio; }; + bool ensure_backend_pair(SDBackendManager& backends, SDBackendModule module); + bool build_core_runners(const Context& ctx, CoreRunners& runners); bool build_vae_runners(const Context& ctx, const VAEOptions& options, VAERunners& runners); bool build_control_net_runner(const Context& ctx, std::shared_ptr& runner); diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index 4256fee4e2..7e215da5f7 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -1,4 +1,5 @@ #include "request.h" +#include "conditioning/conditioner.hpp" #include #include @@ -6,6 +7,9 @@ #include #include "diffusion_engine.h" +#include "generation.h" +#include "model/vae/vae.hpp" +#include "model_definition.h" #include "runtime/denoiser.hpp" namespace sd::pipeline { @@ -43,15 +47,8 @@ namespace sd::pipeline { } enum sample_method_t default_sample_method(const StableDiffusionGGML* sd) { - if (sd != nullptr) { - if (sd_version_is_pid(sd->version)) { - return LCM_SAMPLE_METHOD; - } - if (sd_version_is_dit(sd->version)) { - return EULER_SAMPLE_METHOD; - } - } - return EULER_A_SAMPLE_METHOD; + const auto* definition = sd ? find_model_definition(sd->version) : nullptr; + return definition ? definition->sampling.default_sample_method : EULER_A_SAMPLE_METHOD; } enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method) { @@ -65,21 +62,14 @@ namespace sd::pipeline { return LCM_SCHEDULER; } else if (sample_method == DDIM_TRAILING_SAMPLE_METHOD) { return SIMPLE_SCHEDULER; - } else if (sd != nullptr && (sd_version_is_flux(sd->version) || sd->version == VERSION_QWEN_IMAGE_2_1)) { - return FLUX_SCHEDULER; - } else if (sd != nullptr && sd_version_is_flux2(sd->version)) { - return FLUX2_SCHEDULER; - } else if (sd != nullptr && sd_version_is_ltxav(sd->version)) { - return LTX2_SCHEDULER; - } else if (sd != nullptr && sd_version_is_ideogram4(sd->version)) { - return LOGIT_NORMAL_SCHEDULER; - } else if (sd != nullptr && sd_version_is_llada_image(sd->version)) { - return LLADA_IMAGE_SCHEDULER; + } + if (const auto* definition = sd ? find_model_definition(sd->version) : nullptr) { + return definition->sampling.default_scheduler; } return DISCRETE_SCHEDULER; } - static int64_t resolve_seed(int64_t seed) { + int64_t resolve_seed(int64_t seed) { if (seed >= 0) { return seed; } @@ -140,6 +130,7 @@ namespace sd::pipeline { seed = sd_img_gen_params->seed; batch_count = sd_img_gen_params->batch_count; qwen_image_layers = std::max(0, sd_img_gen_params->qwen_image_layers); + is_layered_image = sd->version == VERSION_QWEN_IMAGE_LAYERED; clip_skip = sd_img_gen_params->clip_skip; shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep; strength = sd_img_gen_params->strength; @@ -150,8 +141,18 @@ namespace sd::pipeline { pm_params = sd_img_gen_params->pm_params; pulid_params = sd_img_gen_params->pulid_params; hires = sd_img_gen_params->hires; + vae_tiling_params = sd_img_gen_params->vae_tiling_params; cache_params = &sd_img_gen_params->cache; + sample_params = sd_img_gen_params->sample_params; resolve(sd); + if (!sd_img_gen_params->control_image.data) + models.erase("control_net"); + if (!sd_img_gen_params->ip_adapter_image.data) + models.erase("ip_adapter"); + if (!pm_params.id_images || pm_params.id_images_count <= 0) + models.erase("photo_maker"); + if (!pulid_params.id_embedding_path || !*pulid_params.id_embedding_path) + models.erase("pulid_weights"); } GenerationRequest::GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params) { @@ -179,12 +180,25 @@ namespace sd::pipeline { guidance = sd_vid_gen_params->sample_params.guidance; high_noise_guidance = sd_vid_gen_params->high_noise_sample_params.guidance; hires = sd_vid_gen_params->hires; + vae_tiling_params = sd_vid_gen_params->vae_tiling_params; + sample_params = sd_vid_gen_params->sample_params; + high_noise_sample_params = sd_vid_gen_params->high_noise_sample_params; + use_high_noise = sd->high_noise_diffusion_model != nullptr; + use_vace = sd->diffusion_model && (sd->diffusion_model->get_desc() == "Wan2.1-VACE-1.3B" || + sd->diffusion_model->get_desc() == "Wan2.x-VACE-14B"); + shifted_timestep = sample_params.shifted_timestep; + if (sd->config_ && sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) + frames = std::min(requested_frames, 32); resolve(sd); + models.erase("control_net"); + models.erase("ip_adapter"); + models.erase("photo_maker"); + models.erase("pulid_weights"); if (frames != requested_frames) { LOG_WARN("align video frames from %d to %d for %s", requested_frames, frames, - model_version_to_str[sd->version]); + sd_version_name(sd->version)); } } @@ -214,6 +228,10 @@ namespace sd::pipeline { spatial_multiple); } + int64_t GenerationRequest::seed_for_batch(int index) const { + return static_cast((static_cast(seed) + index) & INT64_MAX); + } + void GenerationRequest::resolve_hires() { if (!hires.enabled) { return; @@ -329,9 +347,35 @@ namespace sd::pipeline { } void GenerationRequest::resolve(StableDiffusionGGML* sd) { + if (sd->config_) { + models = sd->config_->models; + if (!use_high_noise) + models.erase("high_noise_diffusion_model"); + if (!sd->clip_vision) + models.erase("clip_vision"); + if (!sd->audio_vae_model || frames <= 0) + models.erase("audio_vae"); + if (!sd->audio_encoder || frames <= 0) + models.erase("audio_encoder"); + if (!sd->config_->animatediff_loaded || !sd_version_supports_animatediff(sd->version) || frames <= 0) + models.erase("motion_module"); + if (sd->uses_tae() && !sd->tae_preview_only()) + models.erase("vae"); + else + models.erase("taesd"); + } + if (sd_version_is_minit2i(sd->version) && (width != 512 || height != 512)) { + width = 512; + height = 512; + LOG_WARN("align to fixed size 512x512 for MiniT2I"); + } align_generation_request_size(); resolve_hires(); + if (sd->first_stage_model) + vae_tiling_params = sd->first_stage_model->resolve_tiling_params(vae_tiling_params); seed = resolve_seed(seed); + if (sd->cond_stage_model) + clip_skip = sd->cond_stage_model->resolve_clip_skip(clip_skip); resolve_guidance(sd, &guidance, &use_uncond, &use_img_uncond, has_ref_images); if (sd->high_noise_diffusion_model) { @@ -347,33 +391,44 @@ namespace sd::pipeline { LOG_WARN("timestep shifting is only supported for SDXL models!"); shifted_timestep = 0; } + sample_params.guidance = guidance; + sample_params.shifted_timestep = shifted_timestep; + sample_params.sample_method = resolve_sample_method(sd, sample_params.sample_method); + sample_params.scheduler = resolve_scheduler(sd, sample_params.scheduler, sample_params.sample_method); + sample_params.eta = resolve_eta(sd, sample_params.eta, sample_params.sample_method); + if (frames > 0) { + high_noise_sample_params.guidance = high_noise_guidance; + high_noise_sample_params.sample_method = resolve_sample_method(sd, high_noise_sample_params.sample_method); + high_noise_sample_params.scheduler = resolve_scheduler(sd, high_noise_sample_params.scheduler, high_noise_sample_params.sample_method); + high_noise_sample_params.eta = resolve_eta(sd, high_noise_sample_params.eta, high_noise_sample_params.sample_method); + } } SamplePlan::SamplePlan(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, const GenerationRequest& request) { - sample_method = sd_img_gen_params->sample_params.sample_method; - extra_sample_args = sd_img_gen_params->sample_params.extra_sample_args; - eta = sd_img_gen_params->sample_params.eta; - sample_steps = sd_img_gen_params->sample_params.sample_steps; - resolve(sd, &request, &sd_img_gen_params->sample_params); + sample_method = request.sample_params.sample_method; + extra_sample_args = request.sample_params.extra_sample_args; + eta = request.sample_params.eta; + sample_steps = request.sample_params.sample_steps; + resolve(sd, &request, &request.sample_params); } SamplePlan::SamplePlan(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request) { - sample_method = sd_vid_gen_params->sample_params.sample_method; - extra_sample_args = sd_vid_gen_params->sample_params.extra_sample_args; - eta = sd_vid_gen_params->sample_params.eta; - sample_steps = sd_vid_gen_params->sample_params.sample_steps; + sample_method = request.sample_params.sample_method; + extra_sample_args = request.sample_params.extra_sample_args; + eta = request.sample_params.eta; + sample_steps = request.sample_params.sample_steps; if (sd->high_noise_diffusion_model) { - high_noise_sample_steps = sd_vid_gen_params->high_noise_sample_params.sample_steps; - high_noise_sample_method = sd_vid_gen_params->high_noise_sample_params.sample_method; - high_noise_extra_sample_args = sd_vid_gen_params->high_noise_sample_params.extra_sample_args; - high_noise_eta = sd_vid_gen_params->high_noise_sample_params.eta; + high_noise_sample_steps = request.high_noise_sample_params.sample_steps; + high_noise_sample_method = request.high_noise_sample_params.sample_method; + high_noise_extra_sample_args = request.high_noise_sample_params.extra_sample_args; + high_noise_eta = request.high_noise_sample_params.eta; } moe_boundary = sd_vid_gen_params->moe_boundary; - resolve(sd, &request, &sd_vid_gen_params->sample_params); + resolve(sd, &request, &request.sample_params); } void SamplePlan::resolve(StableDiffusionGGML* sd, @@ -399,10 +454,10 @@ namespace sd::pipeline { } } } else { - scheduler_t scheduler = resolve_scheduler(sd, - sample_params->scheduler, - sample_method); - int sample_seq_len = sd->get_image_seq_len(request->height, request->width); + scheduler = resolve_scheduler(sd, + sample_params->scheduler, + sample_method); + int sample_seq_len = sd->get_image_seq_len(request->height, request->width); if (sd_version_is_ltxav(sd->version) && request->frames > 0) { int latent_frames = ((request->frames - 1) / 8) + 1; sample_seq_len *= latent_frames; @@ -473,14 +528,10 @@ namespace sd::pipeline { scheduler, sd->version, sample_params.extra_sample_args); - size_t t_enc = static_cast(scheduler_steps * hires.denoising_strength); - if (t_enc >= static_cast(scheduler_steps)) { - t_enc = static_cast(scheduler_steps) - 1; - } if (scheduler_steps_out != nullptr) { *scheduler_steps_out = scheduler_steps; } - return std::vector(sigmas.begin() + scheduler_steps - static_cast(t_enc) - 1, + return std::vector(sigmas.begin() + std::max(0, scheduler_steps - effective_steps), sigmas.end()); } diff --git a/src/pipeline/request.h b/src/pipeline/request.h index fac4aba55f..1cf1cee6a3 100644 --- a/src/pipeline/request.h +++ b/src/pipeline/request.h @@ -1,6 +1,7 @@ #ifndef __SD_PIPELINE_REQUEST_H__ #define __SD_PIPELINE_REQUEST_H__ +#include #include #include @@ -16,6 +17,8 @@ namespace sd::pipeline { enum scheduler_t default_scheduler(const StableDiffusionGGML* sd, enum sample_method_t sample_method); + int64_t resolve_seed(int64_t seed); + float resolve_eta(StableDiffusionGGML* sd, float eta, enum sample_method_t sample_method); @@ -23,38 +26,47 @@ namespace sd::pipeline { struct GenerationRequest { std::string prompt; std::string negative_prompt; - int width = -1; - int height = -1; - int clip_skip = -1; - int vae_scale_factor = -1; - int diffusion_model_down_factor = -1; - int64_t seed = -1; - bool use_uncond = false; - bool use_img_uncond = false; - bool use_high_noise_uncond = false; - bool use_high_noise_img_uncond = false; - bool has_ref_images = false; - const sd_cache_params_t* cache_params = nullptr; - int batch_count = 1; - int qwen_image_layers = 3; - int shifted_timestep = 0; - float strength = 1.f; - float control_strength = 0.f; - float eta = 0.f; - sd_guidance_params_t guidance = {}; - sd_guidance_params_t high_noise_guidance = {}; - sd_pm_params_t pm_params = {}; - sd_pulid_params_t pulid_params = {}; - sd_hires_params_t hires = {}; - int frames = -1; - int requested_frames = -1; - int fps = 16; - float vace_strength = 1.f; + std::map models; + int width = -1; + int height = -1; + int clip_skip = -1; + int vae_scale_factor = -1; + int diffusion_model_down_factor = -1; + int64_t seed = -1; + bool use_uncond = false; + bool use_img_uncond = false; + bool use_high_noise_uncond = false; + bool use_high_noise_img_uncond = false; + bool has_ref_images = false; + bool is_layered_image = false; + bool use_high_noise = false; + bool use_vace = false; + const sd_cache_params_t* cache_params = nullptr; + int batch_count = 1; + int qwen_image_layers = 3; + int shifted_timestep = 0; + float strength = 1.f; + float control_strength = 0.f; + float eta = 0.f; + sd_sample_params_t sample_params = {}; + sd_sample_params_t high_noise_sample_params = {}; + sd_guidance_params_t guidance = {}; + sd_guidance_params_t high_noise_guidance = {}; + sd_pm_params_t pm_params = {}; + sd_pulid_params_t pulid_params = {}; + sd_hires_params_t hires = {}; + sd_tiling_params_t vae_tiling_params = {}; + int frames = -1; + int requested_frames = -1; + int fps = 16; + float vace_strength = 1.f; GenerationRequest(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params); GenerationRequest(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params); + int64_t seed_for_batch(int index) const; + void align_generation_request_size(); void align_image_size(int* target_width, int* target_height, const char* label); @@ -81,6 +93,7 @@ namespace sd::pipeline { int sample_steps = 0; int high_noise_sample_steps = 0; int total_steps = 0; + scheduler_t scheduler = SCHEDULER_COUNT; float moe_boundary = 0.f; std::vector sigmas; diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index 2fcff4c9d4..8c911b143a 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -1,15 +1,18 @@ #include "generation.h" +#include "generation_info.h" #include #include #include #include +#include #include "conditioning/wan_audio.h" #include "core/rng.hpp" #include "core/rng_philox.hpp" #include "diffusion_engine.h" #include "model/diffusion/minimax_h3.hpp" +#include "model/upscaler/h3_latent_upscaler.hpp" #include "model/upscaler/ltx_latent_upscaler.hpp" #include "model/vae/audio_vae.hpp" #include "model/vae/vae.hpp" @@ -461,9 +464,35 @@ namespace sd::pipeline { return window; } + // ComfyUI-compatible reference sizing: per-axis rounding to 32 px keeps the reference + // divisible by VAE 16 x patch 2 without padding in the H3 graph. Both modes round rather + // than floor: flooring would under-size in-limit references and diverge from the + // upstream node, which rounds in both branches. + static std::pair minimax_h3_reference_size(int source_w, + int source_h, + int target_w, + int target_h, + RefImageSizeMode mode) { + constexpr double kCanvasMultiple = 32.0; + constexpr double kRefImageShortEdge = 2048.0; + double scale = 1.0; + if (mode == RefImageSizeMode::MAX) { + scale = std::min(1.0, kRefImageShortEdge / std::min(source_w, source_h)); + } else { + double source_area = static_cast(source_w) * source_h; + double target_area = static_cast(target_w) * target_h; + scale = std::min(1.0, std::sqrt(target_area / source_area)); + } + auto align32 = [kCanvasMultiple](double value) { + return std::max(32, static_cast(std::round(value / kCanvasMultiple)) * 32); + }; + return {align32(source_w * scale), align32(source_h * scale)}; + } + static std::optional prepare_video_generation_latents(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, - GenerationRequest* request) { + GenerationRequest* request, + const RefImageParams& ref_image_params) { ImageGenerationLatents latents; int64_t prepare_start_ms = ggml_time_ms(); @@ -576,6 +605,7 @@ namespace sd::pipeline { sd_vid_gen_params->ref_videos_count, sd_vid_gen_params->ref_audios_count); + int64_t ref_image_tokens = 0; for (int i = 0; i < sd_vid_gen_params->ref_images_count; ++i) { auto image = ensure_image_tensor_channels( sd_image_to_tensor(sd_vid_gen_params->ref_images[i]), @@ -584,15 +614,18 @@ namespace sd::pipeline { LOG_ERROR("failed to load MiniMax-H3 reference image %d", i + 1); return std::nullopt; } - int source_w = static_cast(image.shape()[0]); - int source_h = static_cast(image.shape()[1]); - double source_area = static_cast(source_w) * source_h; - double target_area = static_cast(request->width) * request->height; - double scale = std::min(1.0, std::sqrt(target_area / source_area)); - int width = std::max(32, static_cast(std::round(source_w * scale / 32.f)) * 32); - int height = std::max(32, static_cast(std::round(source_h * scale / 32.f)) * 32); - image = resize_reference(image, width, height); - auto latent = sd->encode_first_stage(image); + int source_w = static_cast(image.shape()[0]); + int source_h = static_cast(image.shape()[1]); + auto [width, height] = + minimax_h3_reference_size(source_w, source_h, request->width, request->height, ref_image_params.ref_image_size); + LOG_INFO("MiniMax-H3 reference image %d: %dx%d -> %dx%d", + i + 1, + source_w, + source_h, + width, + height); + image = resize_reference(image, width, height); + auto latent = sd->encode_first_stage(image); if (latent.empty()) { LOG_ERROR("failed to encode MiniMax-H3 reference image %d", i + 1); return std::nullopt; @@ -602,6 +635,10 @@ namespace sd::pipeline { latents.minimax_reference_blocks.push_back({MiniMaxH3ReferenceKind::IMAGE, video_index, -1}); + // The DiT patchifies latents 2 x 2, so the token count is a quarter of the + // latent grid. + ref_image_tokens += static_cast(latents.ref_latents.back().shape()[0] / 2) * + (latents.ref_latents.back().shape()[1] / 2); MiniMaxH3PresentationItem item; item.kind = MiniMaxH3PresentationKind::IMAGE; item.frames.push_back(std::move(image)); @@ -654,8 +691,8 @@ namespace sd::pipeline { int source_index = std::min(reference.frame_count - 1, static_cast(std::floor(frame * source_fps / 24.0))); auto source = ensure_image_tensor_channels( - sd_image_to_tensor(reference.frames[source_index]), - 3); + sd_image_to_tensor(reference.frames[source_index]), + 3); if (source.empty()) { LOG_ERROR("invalid frame %d in MiniMax-H3 reference video %d", source_index + 1, @@ -711,6 +748,18 @@ namespace sd::pipeline { -1, encoded_index}); } + + if (ref_image_tokens > 0) { + int64_t target_tokens = sd->get_image_seq_len(request->width, request->height) / 4 * + sd->video_frames_to_latent_frames(request->frames); + if (target_tokens > 0 && ref_image_tokens > target_tokens * 4) { + LOG_WARN("MiniMax-H3 reference images add %" PRId64 + " tokens versus %" PRId64 + " for the generated video; sampling cost scales with every attended token", + ref_image_tokens, + target_tokens); + } + } } if (!has_references && (!start_image.empty() || !end_image.empty())) { @@ -771,7 +820,7 @@ namespace sd::pipeline { keyframes.shape()[2], 1, 1}, - conditioned_mask); + conditioned_mask); latents.denoise_mask = sd::ops::concat(latents.denoise_mask, keyframe_mask, 2); latents.video_positions = build_ltxv_video_positions(latents.init_latent.shape()[0], latents.init_latent.shape()[1], @@ -1059,10 +1108,10 @@ namespace sd::pipeline { inactive.shape()[1], inactive.shape()[2]}); // [t, h/vae_scale_factor, vae_scale_factor, w/vae_scale_factor, vae_scale_factor] mask_context = mask_context.permute({1, 3, 4, 0, 2}) // [vae_scale_factor, vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] - .reshape({inactive.shape()[0], - inactive.shape()[1], - inactive.shape()[2], - request->vae_scale_factor * request->vae_scale_factor}); // [vae_scale_factor*vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] + .reshape({inactive.shape()[0], + inactive.shape()[1], + inactive.shape()[2], + request->vae_scale_factor * request->vae_scale_factor}); // [vae_scale_factor*vae_scale_factor, t, h/vae_scale_factor, w/vae_scale_factor] if (!ref_image_latent.empty()) { vace_context = sd::ops::concat(ref_image_latent, vace_context, 2); // [b, 2*c, t+1, h/vae_scale_factor, w/vae_scale_factor] @@ -1147,7 +1196,8 @@ namespace sd::pipeline { static std::optional prepare_video_generation_embeds(StableDiffusionGGML* sd, const sd_vid_gen_params_t* sd_vid_gen_params, const GenerationRequest& request, - const ImageGenerationLatents& latents) { + const ImageGenerationLatents& latents, + const RefImageParams& ref_image_params) { ConditionerRunnerEndOnExit conditioner_runner_end{sd->cond_stage_model.get()}; ImageGenerationEmbeds embeds; @@ -1157,9 +1207,7 @@ namespace sd::pipeline { condition_params.zero_out_masked = true; condition_params.ref_images = &latents.ref_images; condition_params.minimax_h3_references = &latents.minimax_presentation_refs; - if (sd_version_is_lingbot_video(sd->version) || sd_version_is_minimax_h3(sd->version)) { - condition_params.ref_image_params.vlm_resize_mode = RefImageResizeMode::AREA; - } + condition_params.ref_image_params = ref_image_params; int64_t prepare_start_ms = ggml_time_ms(); embeds.cond = sd->get_learned_condition(condition_params); @@ -1270,6 +1318,230 @@ namespace sd::pipeline { return result_images; } + static bool hires_upscaler_is_latent_interpolation(enum sd_hires_upscaler_t upscaler) { + switch (upscaler) { + case SD_HIRES_UPSCALER_LATENT: + case SD_HIRES_UPSCALER_LATENT_NEAREST: + case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT: + case SD_HIRES_UPSCALER_LATENT_ANTIALIASED: + case SD_HIRES_UPSCALER_LATENT_BICUBIC: + case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED: + return true; + default: + return false; + } + } + + static sd::ops::InterpolateMode hires_latent_interpolate_mode(enum sd_hires_upscaler_t upscaler, + bool* antialias) { + if (antialias != nullptr) { + *antialias = false; + } + switch (upscaler) { + case SD_HIRES_UPSCALER_LATENT_NEAREST: + return sd::ops::InterpolateMode::Nearest; + case SD_HIRES_UPSCALER_LATENT_NEAREST_EXACT: + return sd::ops::InterpolateMode::NearestExact; + case SD_HIRES_UPSCALER_LATENT_BICUBIC: + return sd::ops::InterpolateMode::Bicubic; + case SD_HIRES_UPSCALER_LATENT_BICUBIC_ANTIALIASED: + if (antialias != nullptr) { + *antialias = true; + } + return sd::ops::InterpolateMode::Bicubic; + case SD_HIRES_UPSCALER_LATENT_ANTIALIASED: + if (antialias != nullptr) { + *antialias = true; + } + return sd::ops::InterpolateMode::Bilinear; + default: + return sd::ops::InterpolateMode::Bilinear; + } + } + + static sd::Tensor upscale_spatial_video_latent(StableDiffusionGGML* sd, + const GenerationRequest& request, + const sd::Tensor& packed_latent, + int audio_length) { + if (sd == nullptr || packed_latent.empty()) { + return {}; + } + if (packed_latent.dim() != 4 && packed_latent.dim() != 5) { + LOG_ERROR("video latent spatial upscale expects a 4D or 5D latent"); + return {}; + } + int64_t target_width = request.hires.target_width / request.vae_scale_factor; + int64_t target_height = request.hires.target_height / request.vae_scale_factor; + if (target_width <= 0 || target_height <= 0) { + LOG_ERROR("video latent spatial upscale target must be positive: %dx%d", + request.hires.target_width, + request.hires.target_height); + return {}; + } + if (packed_latent.shape()[0] == target_width && packed_latent.shape()[1] == target_height) { + LOG_DEBUG("video latent spatial upscale: target equals the draft latent size, skipping"); + return packed_latent; + } + + int latent_channels = sd->get_latent_channel(); + sd::Tensor video_latent = packed_latent; + sd::Tensor audio_latent; + if (packed_latent.shape()[3] > latent_channels) { + video_latent = sd::ops::slice(packed_latent, 3, 0, latent_channels); + audio_latent = sd_version_is_minimax_h3(sd->version) + ? unpack_minimax_h3_audio_latent(packed_latent, audio_length, latent_channels) + : unpack_ltxav_audio_latent(packed_latent, audio_length, latent_channels); + if (audio_latent.empty()) { + LOG_ERROR("failed to unpack audio latent before spatial latent upscale"); + return {}; + } + } + + bool antialias = false; + auto mode = hires_latent_interpolate_mode(request.hires.upscaler, &antialias); + std::vector target_shape = video_latent.shape(); + target_shape[0] = target_width; + target_shape[1] = target_height; + + LOG_INFO("video latent %s spatial upscale: %" PRId64 "x%" PRId64 "x%" PRId64 "x%" PRId64 " -> %" PRId64 "x%" PRId64 "x%" PRId64 "x%" PRId64, + sd_hires_upscaler_name(request.hires.upscaler), + (int64_t)video_latent.shape()[0], + (int64_t)video_latent.shape()[1], + (int64_t)video_latent.shape()[2], + (int64_t)video_latent.shape()[3], + target_shape[0], + target_shape[1], + target_shape[2], + target_shape[3]); + + sd::Tensor upscaled = sd::ops::interpolate(video_latent, target_shape, mode, false, antialias); + if (upscaled.empty()) { + LOG_ERROR("video latent spatial upscale failed"); + return {}; + } + if (audio_latent.empty()) { + return upscaled; + } + return pack_ltxav_audio_and_video_latents(upscaled, audio_latent); + } + + sd::Tensor upscale_h3_learned_video_latent(StableDiffusionGGML* sd, + const char* model_path, + const sd::Tensor& packed_latent, + int audio_length, + int64_t target_width, + int64_t target_height) { + if (sd == nullptr || sd->model_manager == nullptr || packed_latent.empty()) { + return {}; + } + if (strlen(SAFE_STR(model_path)) == 0) { + LOG_ERROR("H3 learned latent upscale requires a hires model path"); + return {}; + } + // The learned network is not an identity mapping at scale 1. + if (packed_latent.shape()[0] == target_width && packed_latent.shape()[1] == target_height) { + LOG_DEBUG("H3 learned latent upscale: target equals the draft latent size, skipping"); + return packed_latent; + } + if (!sd->ensure_backend_pair(SDBackendModule::UPSCALER)) { + return {}; + } + + int latent_channels = sd->get_latent_channel(); + sd::Tensor video_latent = packed_latent; + sd::Tensor audio_latent; + if (packed_latent.shape()[3] > latent_channels) { + video_latent = sd::ops::slice(packed_latent, 3, 0, latent_channels); + audio_latent = unpack_minimax_h3_audio_latent(packed_latent, audio_length, latent_channels); + if (audio_latent.empty()) { + LOG_ERROR("failed to unpack H3 audio latent before learned latent upscale"); + return {}; + } + } + + LOG_INFO("H3 learned latent spatial upscale: latent %dx%dx%dx%d -> %" PRId64 "x%" PRId64 " model output", + (int)video_latent.shape()[0], + (int)video_latent.shape()[1], + (int)video_latent.shape()[2], + (int)video_latent.shape()[3], + target_width, + target_height); + + if (sd->first_stage_model) { + sd->first_stage_model->runner_end(); + } + + auto model_manager = sd->model_manager; + struct UpsamplerScope { + ModelManager& manager; + ModelLoader::FileId owned_source = 0; + std::unique_ptr runner; + std::vector params; + + ~UpsamplerScope() { + if (runner) { + runner->runner_end(); + } + GGML_ASSERT(manager.unregister_param_tensors(params)); + if (owned_source != 0) { + GGML_ASSERT(manager.del_file(owned_source)); + } + } + } scope{*model_manager}; + + const std::string prefix = "h3_latent_upscaler"; + ModelLoader candidate = model_manager->loader(); + ModelLoader::FileId source_file = 0; + if (!candidate.add_file(model_path, prefix + ".", &source_file)) { + LOG_ERROR("init H3 latent upsampler model loader from file failed: '%s'", model_path); + return {}; + } + const bool owns_source = model_manager->loader().file_revision(source_file) == 0; + if (!model_manager->set_loader(std::move(candidate))) { + return {}; + } + scope.owned_source = owns_source ? source_file : 0; + + auto& upsampler = scope.runner; + upsampler = std::make_unique(sd->backend_for(SDBackendModule::UPSCALER), + model_manager->loader().get_tensor_storage_map(), + prefix, + model_manager); + const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); + upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); + if (upsampler->model == nullptr) { + LOG_ERROR("init H3 latent upsampler from metadata failed"); + return {}; + } + + std::map tensors; + upsampler->get_param_tensors(tensors); + for (const auto& entry : tensors) { + scope.params.push_back(entry.second); + } + if (!model_manager->register_param_tensors(ModelComponent::LatentUpsampler, + std::move(tensors), + ModelManager::ResidencyMode::ParamBackend, + sd->backend_for(SDBackendModule::UPSCALER), + sd->params_backend_for(SDBackendModule::UPSCALER)) || + !model_manager->validate_registered_tensors()) { + LOG_ERROR("register H3 latent upsampler tensors with model manager failed"); + return {}; + } + + sd::Tensor upscaled = upsampler->compute(sd->n_threads, video_latent, target_width, target_height); + upsampler->runner_end(); + if (upscaled.empty()) { + LOG_ERROR("H3 learned latent spatial upscale failed"); + return {}; + } + + if (!audio_latent.empty()) { + upscaled = pack_ltxav_audio_and_video_latents(upscaled, audio_latent); + } + return upscaled; + } + sd::Tensor upscale_ltx_spatial_video_latent(StableDiffusionGGML* sd, const char* model_path, const sd::Tensor& packed_latent, @@ -1341,9 +1613,9 @@ namespace sd::pipeline { auto& upsampler = scope.runner; upsampler = std::make_unique(sd->backend_for(SDBackendModule::UPSCALER), - model_manager->loader().get_tensor_storage_map(), - prefix, - model_manager); + model_manager->loader().get_tensor_storage_map(), + prefix, + model_manager); const size_t max_graph_vram_bytes = sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER); upsampler->set_max_graph_vram_bytes(max_graph_vram_bytes); if (upsampler->model == nullptr) { @@ -1526,9 +1798,16 @@ namespace sd::pipeline { img_gen_params.circular_x = sd_vid_gen_params->circular_x; img_gen_params.circular_y = sd_vid_gen_params->circular_y; img_gen_params.image_preprocess = sd_vid_gen_params->image_preprocess; + img_gen_params.ref_image_args = sd_vid_gen_params->ref_image_args; sd->animatediff_num_frames = n_frames; - bool ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); + bool ok = false; + try { + ok = generate_image(sd, &img_gen_params, frames_out, num_frames_out); + } catch (...) { + sd->animatediff_num_frames = 0; + throw; + } sd->animatediff_num_frames = 0; return ok; } @@ -1538,24 +1817,33 @@ namespace sd::pipeline { sd_image_t** frames_out, int* num_frames_out, sd_audio_t** audio_out, - int* fps_out) { + int* fps_out, + nlohmann::json* info_out) { + if (info_out) + *info_out = nlohmann::json(); + GenerationRequest request(sd, sd_vid_gen_params); + const auto info = sd::generation_info(request, *sd_vid_gen_params); if (sd->config_->animatediff_loaded && sd_version_supports_animatediff(sd->version)) { LOG_INFO("AnimateDiff dispatch: %d frames, %dx%d", sd_vid_gen_params->video_frames, sd_vid_gen_params->width, sd_vid_gen_params->height); if (fps_out != nullptr) { *fps_out = std::max(1, sd_vid_gen_params->fps); } - return generate_animatediff_video(sd, sd_vid_gen_params, frames_out, num_frames_out); + auto params = *sd_vid_gen_params; + params.seed = request.seed; + const bool ok = generate_animatediff_video(sd, ¶ms, frames_out, num_frames_out); + if (ok && info_out) + *info_out = info; + return ok; } sd->reset_cancel_flag(); - const RefImageParams ref_image_params; + const RefImageParams ref_image_params = sd->resolve_ref_image_params(sd_vid_gen_params->ref_image_args); int64_t t0 = ggml_time_ms(); sd->vae_tiling_params = sd_vid_gen_params->vae_tiling_params; sd->apply_circular_axes(sd_vid_gen_params->circular_x, sd_vid_gen_params->circular_y); - GenerationRequest request(sd, sd_vid_gen_params); sd::ImagePreprocessor preprocessing(sd_vid_gen_params->image_preprocess.rules); sd_vid_gen_params_t processed_params = *sd_vid_gen_params; if (!preprocessing.prepare_inputs(processed_params, request.width, request.height)) @@ -1567,17 +1855,41 @@ namespace sd::pipeline { bool latent_upscale_enabled = request.hires.enabled; GenerationRequest hires_request = request; if (latent_upscale_enabled) { - if (!sd_version_is_ltxav(sd->version)) { - LOG_ERROR("LTX latent spatial upscale is only supported for LTX video models"); + const bool is_ltxav = sd_version_is_ltxav(sd->version); + const bool is_h3 = sd_version_is_minimax_h3(sd->version); + if (!is_ltxav && !is_h3) { + LOG_ERROR("video latent spatial upscale is only supported for LTX and MiniMax-H3 video models"); return false; } - if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) { - LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL"); - return false; - } - if (strlen(SAFE_STR(request.hires.model_path)) == 0) { - LOG_ERROR("LTX latent spatial upscale is enabled but hires model path was not provided"); - return false; + if (is_ltxav) { + if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) { + LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL"); + return false; + } + if (strlen(SAFE_STR(request.hires.model_path)) == 0) { + LOG_ERROR("LTX latent spatial upscale is enabled but hires model path was not provided"); + return false; + } + } else { + const bool model_upscaler = request.hires.upscaler == SD_HIRES_UPSCALER_MODEL; + if (model_upscaler) { + if (strlen(SAFE_STR(request.hires.model_path)) == 0) { + LOG_ERROR("MiniMax-H3 learned hires upscaler is enabled but hires model path was not provided"); + return false; + } + } else if (!hires_upscaler_is_latent_interpolation(request.hires.upscaler)) { + LOG_ERROR("MiniMax-H3 hires requires an interpolating latent upscaler (e.g. \"Latent\" or \"Latent (bicubic)\") or a MODEL upscaler checkpoint"); + return false; + } + if (sd_vid_gen_params->init_image.data != nullptr || + sd_vid_gen_params->end_image.data != nullptr) { + LOG_ERROR("MiniMax-H3 hires does not support --init-img/--end-img keyframes yet; use T2VA or Ref2VA"); + return false; + } + if (request.hires.target_width <= 0 || request.hires.target_height <= 0) { + LOG_ERROR("MiniMax-H3 hires requires a positive hires target size"); + return false; + } } } @@ -1589,7 +1901,7 @@ namespace sd::pipeline { sd->reset_generation_extensions(); SamplePlan plan(sd, sd_vid_gen_params, request); - auto latent_inputs_opt = prepare_video_generation_latents(sd, sd_vid_gen_params, &request); + auto latent_inputs_opt = prepare_video_generation_latents(sd, sd_vid_gen_params, &request, ref_image_params); if (!latent_inputs_opt.has_value()) { return false; } @@ -1598,13 +1910,14 @@ namespace sd::pipeline { auto embeds_opt = prepare_video_generation_embeds(sd, sd_vid_gen_params, request, - latents); + latents, + ref_image_params); if (!embeds_opt) { return false; } ImageGenerationEmbeds embeds = std::move(*embeds_opt); if (latent_upscale_enabled) { - LOG_INFO("generate_video %dx%dx%d -> LTX latent spatial upscale", + LOG_INFO("generate_video %dx%dx%d -> video latent spatial upscale", request.width, request.height, request.frames); @@ -1686,9 +1999,9 @@ namespace sd::pipeline { embeds.img_uncond, sd::Tensor(), 0.f, - sd_vid_gen_params->sample_params.guidance, + request.guidance, plan.eta, - sd_vid_gen_params->sample_params.shifted_timestep, + request.shifted_timestep, plan.sample_method, sd->is_flow_denoiser(), plan.extra_sample_args, @@ -1716,16 +2029,31 @@ namespace sd::pipeline { LOG_ERROR("cancelling generation before latent upscale"); return false; } - int64_t upscale_start = ggml_time_ms(); - sd::Tensor upscaled_latent = upscale_ltx_spatial_video_latent(sd, - request.hires.model_path, - final_latent, - latents.audio_length); - int64_t upscale_end = ggml_time_ms(); + int64_t upscale_start = ggml_time_ms(); + sd::Tensor upscaled_latent; + if (sd_version_is_ltxav(sd->version)) { + upscaled_latent = upscale_ltx_spatial_video_latent(sd, + request.hires.model_path, + final_latent, + latents.audio_length); + } else if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { + upscaled_latent = upscale_h3_learned_video_latent(sd, + request.hires.model_path, + final_latent, + latents.audio_length, + request.hires.target_width / request.vae_scale_factor, + request.hires.target_height / request.vae_scale_factor); + } else { + upscaled_latent = upscale_spatial_video_latent(sd, + request, + final_latent, + latents.audio_length); + } + int64_t upscale_end = ggml_time_ms(); if (upscaled_latent.empty()) { return false; } - LOG_INFO("LTX latent spatial upscale completed, taking %.2fs", + LOG_INFO("video latent spatial upscale completed, taking %.2fs", (upscale_end - upscale_start) * 1.0f / 1000); x_t = std::move(upscaled_latent); @@ -1734,7 +2062,7 @@ namespace sd::pipeline { int upscaled_latent_frames = static_cast(x_t.shape()[2]); int upscaled_frames = sd->latent_frames_to_video_frames(upscaled_latent_frames); if (upscaled_frames != hires_request.frames) { - LOG_INFO("LTX latent upsampler output latent frames %d, frames %d -> %d", + LOG_INFO("video latent upsampler output latent frames %d, frames %d -> %d", upscaled_latent_frames, hires_request.frames, upscaled_frames); @@ -1765,7 +2093,8 @@ namespace sd::pipeline { latents.audio_length = target_audio_length; } } - if ((request.hires.target_width > 0 || request.hires.target_height > 0) && + if (sd_version_is_ltxav(sd->version) && + (request.hires.target_width > 0 || request.hires.target_height > 0) && (request.hires.target_width != hires_request.width || request.hires.target_height != hires_request.height)) { LOG_WARN("LTX latent spatial upsampler output is %dx%d; ignoring hires target %dx%d", hires_request.width, @@ -1779,7 +2108,8 @@ namespace sd::pipeline { LOG_ERROR("cancelling generation before latent upscale refine"); return false; } - if (!apply_ltxv_refine_image_conditioning(sd, + if (sd_version_is_ltxav(sd->version) && + !apply_ltxv_refine_image_conditioning(sd, sd_vid_gen_params, hires_request, latents, @@ -1788,6 +2118,14 @@ namespace sd::pipeline { &hires_video_positions)) { return false; } + // H3 keyframe conditioning must be re-encoded at the hires resolution. + if (sd_version_is_minimax_h3(sd->version) && + !latents.denoise_mask.empty() && + (latents.denoise_mask.shape()[0] != x_t.shape()[0] || + latents.denoise_mask.shape()[1] != x_t.shape()[1])) { + LOG_ERROR("MiniMax-H3 hires cannot rescale an existing video denoise mask"); + return false; + } noise = sd::Tensor::randn_like(x_t, sd->rng); W = hires_request.width / hires_request.vae_scale_factor; @@ -1798,7 +2136,7 @@ namespace sd::pipeline { std::vector hires_sigma_sched = make_hires_sigma_schedule(sd, request.hires, - sd_vid_gen_params->sample_params, + request.sample_params, hires_sample_method, plan.sample_steps, sd->get_image_seq_len(hires_request.height, hires_request.width) * T, @@ -1808,7 +2146,7 @@ namespace sd::pipeline { hires_sample_method); LOG_VERBOSE("sample(latent upscale) %dx%dx%d", W, H, T); - LOG_INFO("LTX latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s", + LOG_INFO("video latent spatial upscale refine: scheduler_steps=%d, denoising_strength=%.2f, sampler=%s, sigma_sched_size=%zu%s", hires_scheduler_steps, request.hires.denoising_strength, sampling_methods_str[hires_sample_method], @@ -1821,13 +2159,13 @@ namespace sd::pipeline { x_t, std::move(noise), embeds.cond, - hires_request.use_uncond ? embeds.uncond : SDCondition(), + hires_request.use_uncond ? embeds.uncond : SDCondition(), embeds.img_uncond, sd::Tensor(), 0.f, - sd_vid_gen_params->sample_params.guidance, + request.guidance, hires_eta, - sd_vid_gen_params->sample_params.shifted_timestep, + request.shifted_timestep, hires_sample_method, sd->is_flow_denoiser(), plan.extra_sample_args, @@ -1959,6 +2297,14 @@ namespace sd::pipeline { generated_audio->sample_count = want_samples; } } + if (info_out) { + try { + *info_out = info; + } catch (...) { + free_sd_audio(generated_audio); + throw; + } + } if (audio_out != nullptr) { *audio_out = generated_audio; } else { diff --git a/src/runtime/cache_dit.hpp b/src/runtime/cache_dit.hpp index 273aa0ef3a..666e61cade 100644 --- a/src/runtime/cache_dit.hpp +++ b/src/runtime/cache_dit.hpp @@ -842,7 +842,7 @@ struct CacheDitConditionState { const float* input_data = input.data(); float diff = CacheDitState::calculate_residual_diff( - it->second.prev_input.data(), input_data, ne); + it->second.prev_input.data(), input_data, ne); float effective_threshold = config.residual_diff_threshold; if (config.Fn_compute_blocks > 0) { diff --git a/src/runtime/guidance.cpp b/src/runtime/guidance.cpp index 0df0f7a772..d02e953af6 100644 --- a/src/runtime/guidance.cpp +++ b/src/runtime/guidance.cpp @@ -192,8 +192,8 @@ namespace sd::guidance { if (has_tensor(input.pred_img_uncond)) { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; output.pred = pred_img_uncond + - image_guidance_scale_ * (pred_uncond - pred_img_uncond) + - guidance_scale * (pred_cond - pred_uncond); + image_guidance_scale_ * (pred_uncond - pred_img_uncond) + + guidance_scale * (pred_cond - pred_uncond); } else { output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond); @@ -252,8 +252,8 @@ namespace sd::guidance { if (has_tensor(input.pred_img_uncond)) { const sd::Tensor& pred_img_uncond = *input.pred_img_uncond; output.pred = pred_img_uncond + - image_guidance_scale_ * (pred_uncond - pred_img_uncond) + - guidance_scale * (pred_cond - pred_uncond); + image_guidance_scale_ * (pred_uncond - pred_img_uncond) + + guidance_scale * (pred_cond - pred_uncond); } else { output.pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond); } diff --git a/src/runtime/image_preprocess.cpp b/src/runtime/image_preprocess.cpp index a8df7c3b56..445f0d1fc8 100644 --- a/src/runtime/image_preprocess.cpp +++ b/src/runtime/image_preprocess.cpp @@ -6,6 +6,42 @@ namespace sd { + std::pair select_image_size_bucket(int source_width, int source_height, int area, int factor, int fallback_width, int fallback_height) { + int width = fallback_width; + int height = fallback_height; + double ratio = static_cast(source_width) / source_height; + double best_crop = INFINITY; + int64_t best_area_error = INT64_MAX; + int min_size = 4 * factor; + for (int w = min_size; w < 2LL * area / min_size; w += factor) { + for (int h = min_size; h < 2LL * area / w; h += factor) { + int64_t pixels = int64_t(w) * h; + if (pixels < 0.9 * area || pixels > 1.1 * area) { + continue; + } + double candidate_ratio = static_cast(w) / h; + double crop = std::round((1.0 - std::min(candidate_ratio / ratio, ratio / candidate_ratio)) * 1e9); + int64_t area_error = std::abs(pixels - area); + if (crop < best_crop || (crop == best_crop && (area_error < best_area_error || (area_error == best_area_error && w > width)))) { + width = w; + height = h; + best_crop = crop; + best_area_error = area_error; + } + } + } + return {width, height}; + } + + Tensor resize_and_center_crop(const Tensor& image, int width, int height, ops::InterpolateMode filter) { + double scale = std::max(static_cast(width) / image.shape()[0], static_cast(height) / image.shape()[1]); + int w = static_cast(std::ceil(image.shape()[0] * scale)); + int h = static_cast(std::ceil(image.shape()[1] * scale)); + auto resized = ops::interpolate(image, {w, h, image.shape()[2], image.shape()[3]}, filter); + resized = ops::slice(resized, 0, (w - width) / 2, (w - width) / 2 + width); + return ops::slice(resized, 1, (h - height) / 2, (h - height) / 2 + height); + } + static constexpr std::pair image_targets[] = { {"init", ImageTarget::Init}, {"end", ImageTarget::End}, @@ -343,8 +379,13 @@ namespace sd { if (target == ImageTarget::Mask && has_init_transform_) { auto options = resolve_options(target, index); if (image.width != init_transform_.source_width || image.height != init_transform_.source_height) { - fail("mask and init source dimensions must match"); - return {}; + ImagePreprocessor alignment; + tensor = alignment.apply_geometry(tensor, target, index, init_transform_.source_width, init_transform_.source_height, + ImageResizeMode::CropResize, ops::InterpolateMode::NearestExact); + if (tensor.empty()) { + fail("could not align mask to init source dimensions"); + return {}; + } } bool geometry_override = options.count("width") || options.count("anchor") || (options.count("mode") && options.at("mode") != "auto"); diff --git a/src/runtime/image_preprocess.h b/src/runtime/image_preprocess.h index b75c4e0e50..7dfc458992 100644 --- a/src/runtime/image_preprocess.h +++ b/src/runtime/image_preprocess.h @@ -3,6 +3,7 @@ #include #include +#include #include #include "core/tensor.hpp" @@ -10,6 +11,9 @@ namespace sd { + std::pair select_image_size_bucket(int source_width, int source_height, int area, int factor, int fallback_width, int fallback_height); + Tensor resize_and_center_crop(const Tensor& image, int width, int height, ops::InterpolateMode filter = ops::InterpolateMode::Lanczos); + enum class ImageTarget { Init, End, diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index d476f92d6b..bca4b79bac 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -1,4 +1,5 @@ #include "stable-diffusion.h" +#include "generation_info.h" #include #include @@ -171,10 +172,13 @@ const char* prediction_to_str[] = { "sefi_flow", "minit2i_flow", "sensenova_u1_flow", + "flux3_flow", + "h3_av_flow", + "ming_image_flow", }; const char* sd_prediction_name(enum prediction_t prediction) { - if (prediction < PREDICTION_COUNT) { + if (prediction >= EPS_PRED && prediction < PREDICTION_COUNT) { return prediction_to_str[prediction]; } return NONE_STR; @@ -633,6 +637,7 @@ void sd_vid_gen_params_init(sd_vid_gen_params_t* sd_vid_gen_params) { sd_vid_gen_params->hires.custom_sigmas_count = 0; sd_vid_gen_params->circular_x = false; sd_vid_gen_params->circular_y = false; + sd_vid_gen_params->ref_image_args = ""; sd_cache_params_init(&sd_vid_gen_params->cache); } @@ -704,13 +709,6 @@ SD_API bool sd_ctx_supports_video_generation(const sd_ctx_t* sd_ctx) { return sd_version_supports_video_generation(sd_ctx->sd->version); } -SD_API bool sd_ctx_supports_ref_image_conditioning(const sd_ctx_t* sd_ctx) { - if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { - return false; - } - return sd_version_supports_ref_image_conditioning(sd_ctx->sd->version); -} - SD_API bool sd_ctx_load_control_net(sd_ctx_t* sd_ctx, const char* path) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr || path == nullptr) { return false; @@ -729,14 +727,20 @@ SD_API bool sd_ctx_has_control_net(const sd_ctx_t* sd_ctx) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr) { return false; } - return sd_ctx->sd->control_net != nullptr; + if (sd_ctx->sd->control_net != nullptr) { + return true; + } + // z-image keeps control net weights inside the diffusion runner, registered + // under ModelComponent::ControlNet instead of a separate control_net runner. + return sd_ctx->sd->model_manager != nullptr && + sd_ctx->sd->model_manager->registered_params_size({ModelComponent::ControlNet}) > 0; } const char* sd_get_model_version_name(const sd_ctx_t* sd_ctx) { if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_ctx->sd->version >= VERSION_COUNT) { return "Unknown"; } - return model_version_to_str[sd_ctx->sd->version]; + return sd_version_name(sd_ctx->sd->version); } enum sample_method_t sd_get_default_sample_method(const sd_ctx_t* sd_ctx) { @@ -749,43 +753,62 @@ enum scheduler_t sd_get_default_scheduler(const sd_ctx_t* sd_ctx, enum sample_me SD_API bool generate_image(sd_ctx_t* sd_ctx, const sd_img_gen_params_t* params, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) - *images_out = nullptr; - if (num_images_out != nullptr) - *num_images_out = 0; - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || params == nullptr) { + sd_image_generation_result_t** result_out) { + if (result_out) + *result_out = nullptr; + if (!result_out || !sd_ctx || !sd_ctx->sd || !params) return false; - } StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); - return execution.ready && sd::pipeline::generate_image(sd_ctx->sd, params, images_out, num_images_out); + if (!execution.ready) + return false; + sd_image_t* images = nullptr; + int count = 0; + std::vector infos; + try { + if (!sd::pipeline::generate_image(sd_ctx->sd, params, &images, &count, &infos)) { + free_sd_images(images, count); + return false; + } + } catch (const std::exception& e) { + LOG_ERROR("image generation failed: %s", e.what()); + free_sd_images(images, count); + return false; + } + return sd::make_image_generation_result(images, count, infos, result_out); } SD_API bool generate_video(sd_ctx_t* sd_ctx, - const sd_vid_gen_params_t* sd_vid_gen_params, - sd_image_t** frames_out, - int* num_frames_out, - sd_audio_t** audio_out, - int* fps_out) { - if (frames_out != nullptr) - *frames_out = nullptr; - if (audio_out != nullptr) - *audio_out = nullptr; - if (num_frames_out != nullptr) - *num_frames_out = 0; - if (fps_out != nullptr) - *fps_out = 0; - if (sd_ctx == nullptr || sd_ctx->sd == nullptr || sd_vid_gen_params == nullptr) { + const sd_vid_gen_params_t* params, + sd_video_generation_result_t** result_out) { + if (result_out) + *result_out = nullptr; + if (!result_out || !sd_ctx || !sd_ctx->sd || !params) return false; - } - StableDiffusionGGML::ExecutionScope execution(*sd_ctx->sd); - if (!execution.ready) { + if (!execution.ready) + return false; + auto* result = static_cast(calloc(1, sizeof(sd_video_generation_result_t))); + if (!result) + return false; + try { + nlohmann::json info; + if (!sd::pipeline::generate_video(sd_ctx->sd, params, &result->frames, &result->frame_count, + &result->audio, &result->fps, &info)) { + free_sd_video_generation_result(result); + return false; + } + result->generation_info = sd::copy_generation_info(info); + if (!result->generation_info) { + free_sd_video_generation_result(result); + return false; + } + } catch (const std::exception& e) { + LOG_ERROR("video generation failed: %s", e.what()); + free_sd_video_generation_result(result); return false; } - - return sd::pipeline::generate_video(sd_ctx->sd, sd_vid_gen_params, frames_out, num_frames_out, audio_out, fps_out); + *result_out = result; + return true; } SD_API void free_sd_images(sd_image_t* result_images, int num_images) { diff --git a/src/upscaler.cpp b/src/upscaler.cpp index dfe5e4b22f..c9bbbcb6b2 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -1,6 +1,7 @@ #include "upscaler.h" #include "core/ggml_extend_backend.h" #include "core/util.h" +#include "generation_info.h" #include "model_loader.h" #include "runtime/tiling.h" #include "stable-diffusion.h" @@ -259,14 +260,11 @@ upscaler_ctx_t* new_upscaler_ctx(const char* esrgan_path_c_str, bool upscale(upscaler_ctx_t* upscaler_ctx, sd_image_t input_image, uint32_t upscale_factor, - sd_image_t** images_out, - int* num_images_out) { - if (images_out != nullptr) { - *images_out = nullptr; - } - if (num_images_out != nullptr) { - *num_images_out = 0; - } + sd_image_generation_result_t** result_out) try { + if (result_out) + *result_out = nullptr; + if (!result_out) + return false; if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) { return false; } @@ -276,21 +274,17 @@ bool upscale(upscaler_ctx_t* upscaler_ctx, return false; } + auto release_images = [](sd_image_t* images) { free_sd_images(images, 1); }; + std::unique_ptr image_owner(result_images, release_images); + const std::vector infos{sd::upscale_generation_info(upscale_factor)}; result_images[0] = upscaler_ctx->upscaler->upscale(input_image, upscale_factor); - if (result_images[0].data == nullptr) { - free(result_images); + if (!result_images[0].data) return false; - } - if (num_images_out != nullptr) { - *num_images_out = 1; - } - if (images_out != nullptr) { - *images_out = result_images; - } else { - free_sd_images(result_images, 1); - } - return true; + return sd::make_image_generation_result(image_owner.release(), 1, infos, result_out); +} catch (const std::exception& e) { + LOG_ERROR("upscale failed: %s", e.what()); + return false; } int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) { @@ -329,6 +323,8 @@ int get_upscaler_model_scale(const char* model_path) { } void free_upscaler_ctx(upscaler_ctx_t* upscaler_ctx) { + if (!upscaler_ctx) + return; if (upscaler_ctx->upscaler != nullptr) { delete upscaler_ctx->upscaler; upscaler_ctx->upscaler = nullptr; diff --git a/thirdparty/cutlass b/thirdparty/cutlass new file mode 160000 index 0000000000..ad7b2f5e84 --- /dev/null +++ b/thirdparty/cutlass @@ -0,0 +1 @@ +Subproject commit ad7b2f5e84fcfa124cb02b91d5bd26d238c0459e