From 45a63ac460b9403df7a739bbc8662cb7d867eace Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 1 Nov 2026 21:18:37 +0800 Subject: [PATCH 01/57] refactor: move model-specific operators and fusions into backend extensions (#2253) --- .gitmodules | 3 + CMakeLists.txt | 42 + cmake/backend_extensions.h.in | 13 + cmake/cutlass.cmake | 20 + cmake/embed_spirv.cmake | 3 + cmake/ggml.cmake | 22 +- cmake/hip.cmake | 12 + docs/build.md | 33 +- docs/ggml_extensions.md | 465 ++++++++++ docs/int8_convrot.md | 13 +- docs/sage_attention.md | 5 +- ggml | 2 +- scripts/format-code.ps1 | 15 +- scripts/format-code.sh | 7 +- src/core/compute_workspace.cpp | 55 +- src/core/cpu/int8.cpp | 181 ++++ src/core/cpu/registration.cpp | 12 + src/core/cpu/registration.h | 14 + src/core/cpu/rope.cpp | 139 +++ src/core/cuda/CMakeLists.txt | 84 ++ src/core/cuda/affine.cu | 183 ++++ src/core/cuda/convrot.cu | 296 +++++++ src/core/cuda/convrot_kernel.h | 251 ++++++ src/core/cuda/cute_attention.cu | 792 +++++++++++++++++ src/core/cuda/i8_gated_residual.cu | 461 ++++++++++ src/core/cuda/i8_gemm.cu | 330 +++++++ src/core/cuda/i8_gemm.h | 9 + src/core/cuda/i8_projection.cu | 394 +++++++++ src/core/cuda/i8_projection.h | 11 + src/core/cuda/int8.cu | 739 ++++++++++++++++ src/core/cuda/registration.cu | 36 + src/core/cuda/registration.h | 13 + src/core/cuda/rope.cu | 337 ++++++++ src/core/cuda/rope.h | 11 + src/core/cuda/runtime.h | 111 +++ src/core/cuda/sage_attention.cu | 387 +++++++++ src/core/cuda/sageattention/LICENSE | 201 +++++ src/core/cuda/sageattention/README.md | 28 + src/core/cuda/sageattention/attn.cuh | 592 +++++++++++++ src/core/cuda/sageattention/attn_fp8.cuh | 622 +++++++++++++ src/core/cuda/sageattention/attn_utils.cuh | 814 ++++++++++++++++++ src/core/cuda/sageattention/cp_async.cuh | 148 ++++ src/core/cuda/sageattention/math.cuh | 166 ++++ src/core/cuda/sageattention/mma.cuh | 683 +++++++++++++++ .../cuda/sageattention/numeric_conversion.cuh | 148 ++++ src/core/cuda/sageattention/permuted_smem.cuh | 201 +++++ src/core/cuda/sol_attention.cu | 409 +++++++++ src/core/ggml_compat.cpp | 185 ++++ src/core/ggml_compat.h | 43 + src/core/ggml_extend.h | 1 + src/core/ggml_extend_backend.cpp | 47 +- src/core/ggml_extensions.cpp | 364 ++++++++ src/core/ggml_extensions.h | 66 ++ src/core/ggml_graph_cut.cpp | 64 +- src/core/vulkan/CMakeLists.txt | 69 ++ src/core/vulkan/int8.cpp | 806 +++++++++++++++++ src/core/vulkan/rope.cpp | 226 +++++ src/core/vulkan/shaders/modulate_convrot.comp | 160 ++++ .../vulkan/shaders/mul_mat_i8_tensorwise.comp | 74 ++ .../shaders/mul_mat_i8_tensorwise_cm1.comp | 106 +++ .../shaders/mul_mat_i8_tensorwise_cm2.comp | 172 ++++ .../vulkan/shaders/quantize_i8_convrot.comp | 105 +++ .../shaders/quantize_i8_convrot_subgroup.comp | 140 +++ src/core/vulkan/shaders/rms_rope_apply.comp | 108 +++ src/core/vulkan/shaders/rope_apply.comp | 66 ++ src/model/common/rope.hpp | 4 +- src/model/diffusion/minimax_h3.hpp | 4 +- thirdparty/cutlass | 1 + 68 files changed, 12173 insertions(+), 151 deletions(-) create mode 100644 cmake/backend_extensions.h.in create mode 100644 cmake/cutlass.cmake create mode 100644 cmake/embed_spirv.cmake create mode 100644 cmake/hip.cmake create mode 100644 docs/ggml_extensions.md create mode 100644 src/core/cpu/int8.cpp create mode 100644 src/core/cpu/registration.cpp create mode 100644 src/core/cpu/registration.h create mode 100644 src/core/cpu/rope.cpp create mode 100644 src/core/cuda/CMakeLists.txt create mode 100644 src/core/cuda/affine.cu create mode 100644 src/core/cuda/convrot.cu create mode 100644 src/core/cuda/convrot_kernel.h create mode 100644 src/core/cuda/cute_attention.cu create mode 100644 src/core/cuda/i8_gated_residual.cu create mode 100644 src/core/cuda/i8_gemm.cu create mode 100644 src/core/cuda/i8_gemm.h create mode 100644 src/core/cuda/i8_projection.cu create mode 100644 src/core/cuda/i8_projection.h create mode 100644 src/core/cuda/int8.cu create mode 100644 src/core/cuda/registration.cu create mode 100644 src/core/cuda/registration.h create mode 100644 src/core/cuda/rope.cu create mode 100644 src/core/cuda/rope.h create mode 100644 src/core/cuda/runtime.h create mode 100644 src/core/cuda/sage_attention.cu create mode 100644 src/core/cuda/sageattention/LICENSE create mode 100644 src/core/cuda/sageattention/README.md create mode 100644 src/core/cuda/sageattention/attn.cuh create mode 100644 src/core/cuda/sageattention/attn_fp8.cuh create mode 100644 src/core/cuda/sageattention/attn_utils.cuh create mode 100644 src/core/cuda/sageattention/cp_async.cuh create mode 100644 src/core/cuda/sageattention/math.cuh create mode 100644 src/core/cuda/sageattention/mma.cuh create mode 100644 src/core/cuda/sageattention/numeric_conversion.cuh create mode 100644 src/core/cuda/sageattention/permuted_smem.cuh create mode 100644 src/core/cuda/sol_attention.cu create mode 100644 src/core/ggml_compat.cpp create mode 100644 src/core/ggml_compat.h create mode 100644 src/core/ggml_extensions.cpp create mode 100644 src/core/ggml_extensions.h create mode 100644 src/core/vulkan/CMakeLists.txt create mode 100644 src/core/vulkan/int8.cpp create mode 100644 src/core/vulkan/rope.cpp create mode 100644 src/core/vulkan/shaders/modulate_convrot.comp create mode 100644 src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp create mode 100644 src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp create mode 100644 src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp create mode 100644 src/core/vulkan/shaders/quantize_i8_convrot.comp create mode 100644 src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp create mode 100644 src/core/vulkan/shaders/rms_rope_apply.comp create mode 100644 src/core/vulkan/shaders/rope_apply.comp create mode 160000 thirdparty/cutlass 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..235556facd 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) + 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 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 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 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) +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/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in new file mode 100644 index 0000000000..cf41592fbe --- /dev/null +++ b/cmake/backend_extensions.h.in @@ -0,0 +1,13 @@ +#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); + +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 39d77eae35..55d0067bcc 100644 --- a/cmake/ggml.cmake +++ b/cmake/ggml.cmake @@ -9,17 +9,19 @@ if(NOT TARGET ggml) add_library(ggml ALIAS ggml::ggml) endif() -get_target_property(sd_ggml_imported ggml IMPORTED) -if(sd_ggml_imported) - set(sd_ggml_private_include "${SD_GGML_SOURCE_DIR}/src") -else() - get_target_property(sd_ggml_private_include ggml SOURCE_DIR) -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}") if(SD_USE_UPSTREAM_GGML) target_compile_definitions(${SD_LIB} PUBLIC SD_USE_UPSTREAM_GGML) 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/docs/build.md b/docs/build.md index 0e7b1a99ef..62966ec6f5 100644 --- a/docs/build.md +++ b/docs/build.md @@ -32,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 @@ -50,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 @@ -95,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. @@ -107,7 +111,7 @@ cmake .. -DSD_CUDA=ON cmake --build . --config Release ``` -CUDA builds with the bundled patched GGML use the CUTLASS submodule by default +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 @@ -117,14 +121,17 @@ To build with the existing CUDA implementations without CUTLASS, configure from the repository root with: ```shell -cmake -S . -B build -DSD_CUDA=ON -DGGML_CUDA_CUTLASS_PATH:PATH= +cmake -S . -B build -DSD_CUDA=ON -DSD_CUDA_CUTLASS_PATH:PATH= cmake --build build --config Release ``` -Alternatively, set `GGML_CUDA_CUTLASS_PATH` to another local CUTLASS source +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 `-UGGML_CUDA_CUTLASS_PATH` when reconfiguring to restore the bundled default. -Non-CUDA builds do not require CUTLASS. +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 diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md new file mode 100644 index 0000000000..6787c443ca --- /dev/null +++ b/docs/ggml_extensions.md @@ -0,0 +1,465 @@ +# 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. + +`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`, `RMSNorm -> Mul(weight) -> Mul(modulation) -> Add -> +Add(shift) -> 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/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/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/ggml b/ggml index c57b71a344..d49c49e1dd 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit c57b71a3448809d07d49360273201bb5a5e7cafb +Subproject commit d49c49e1dd361ba3c0a77c90b5b1404b83969aa8 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/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp index de1cea8631..4e531755b9 100644 --- a/src/core/compute_workspace.cpp +++ b/src/core/compute_workspace.cpp @@ -6,9 +6,9 @@ #include #include +#include "core/ggml_compat.h" #include "core/ggml_extend_backend.h" #include "core/util.h" -#include "ggml-impl.h" namespace sd { ComputeWorkspace::~ComputeWorkspace() { @@ -69,7 +69,7 @@ namespace sd { bool ComputeWorkspace::allocate(ggml_cgraph* graph, const AssignNodes& assign_nodes) { GGML_ASSERT(!active_); - const size_t capacity = static_cast(graph->n_nodes + graph->n_leafs) + 8; + 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); @@ -101,11 +101,11 @@ namespace sd { tensors.push_back(tensor); } }; - for (int i = 0; i < graph->n_nodes; ++i) { - visit(graph->nodes[i]); + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + visit(ggml_graph_node(graph, i)); } - for (int i = 0; i < graph->n_leafs; ++i) { - visit(graph->leafs[i]); + for (int i = 0; i < ggml_graph_n_leafs(graph); ++i) { + visit(ggml_graph_leaf(graph, i)); } for (size_t i = 0; i < tensors.size(); ++i) { visit(tensors[i]->view_src); @@ -120,9 +120,8 @@ namespace sd { const auto* bytes = reinterpret_cast(&value); key.insert(key.end(), bytes, bytes + sizeof(value)); }; - append(graph->n_nodes); - append(graph->n_leafs); - append(graph->order); + append(ggml_graph_n_nodes(graph)); + append(ggml_graph_n_leafs(graph)); for (auto tensor : tensors) { append(tensor); append(tensor->type); @@ -147,11 +146,11 @@ namespace sd { append(usage); append(type); } + } else { + const auto alignment_offset = reinterpret_cast(tensor->data) % ggml_backend_get_alignment(external); + append(alignment_offset); } - const size_t index = ggml_hash_find(&graph->visited_hash_set, tensor); - const int32_t uses = index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, index) - ? graph->use_counts[index] - : 0; + const int32_t uses = ggml_graph_use_count(graph, tensor); append(uses); } // Match the full allocation metadata, including topology and external placement. @@ -187,28 +186,18 @@ 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; - } - } - 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]); - } - copy_graph->order = graph->order; - // Preserve fusion metadata, including uses outside a graph view, without rebuilding its node order. - for (const auto& entry : copies) { - const size_t source_index = ggml_hash_find(&graph->visited_hash_set, entry.first); - if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(graph->visited_hash_set.used, source_index)) { - const size_t copy_index = ggml_hash_insert(©_graph->visited_hash_set, entry.second); - copy_graph->use_counts[copy_index] = graph->use_counts[source_index]; + // 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); + 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; auto scheduler = make_scheduler(graph_size); if (scheduler != nullptr) { 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/registration.cpp b/src/core/cpu/registration.cpp new file mode 100644 index 0000000000..a4f5bd32be --- /dev/null +++ b/src/core/cpu/registration.cpp @@ -0,0 +1,12 @@ +#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); + } + +} +#endif diff --git a/src/core/cpu/registration.h b/src/core/cpu/registration.h new file mode 100644 index 0000000000..c711f6fac9 --- /dev/null +++ b/src/core/cpu/registration.h @@ -0,0 +1,14 @@ +#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_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..7bca6aa33f --- /dev/null +++ b/src/core/cuda/CMakeLists.txt @@ -0,0 +1,84 @@ +cmake_policy(SET CMP0079 NEW) + +function(sd_cuda_sources VENDOR LANGUAGE) + set(SOURCES) + foreach(OP rope int8 convrot 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}/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/convrot.cu b/src/core/cuda/convrot.cu new file mode 100644 index 0000000000..d5d159be14 --- /dev/null +++ b/src/core/cuda/convrot.cu @@ -0,0 +1,296 @@ +#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, + 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 == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && + nodes[2]->op == GGML_OP_MUL && nodes[2]->src[0] == nodes[1] && + nodes[3]->op == GGML_OP_ADD && nodes[3]->src[0] == nodes[1] && nodes[3]->src[1] == nodes[2] && + nodes[4]->op == GGML_OP_ADD && nodes[4]->src[0] == nodes[3] && value == nodes[4]) { + op.kind = RMS_MOD; + op.weight = nodes[1]->src[1]; + op.modulation = nodes[2]->src[1]; + op.shift = nodes[4]->src[1]; + op.eps = param_f32(first, 0); + if (op.k < 1024 || !vector(op.weight, op.k) || !vector(op.modulation, op.k) || !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 ? 6 : 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 = 2; 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 (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) { + auto* add = value->src[0]; + auto* weight = add->src[0]; + if (weight->op == GGML_OP_MUL) { + chain[0] = weight->src[0]; + chain[1] = weight; + chain[2] = add->src[1]; + chain[3] = add; + chain[4] = value; + chain[5] = quant; + count = 6; + } + } 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; + 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 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_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.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..748c08e7de --- /dev/null +++ b/src/core/cuda/convrot_kernel.h @@ -0,0 +1,251 @@ +#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]); + value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), 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/registration.cu b/src/core/cuda/registration.cu new file mode 100644 index 0000000000..d03b114cca --- /dev/null +++ b/src/core/cuda/registration.cu @@ -0,0 +1,36 @@ +#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); +#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_i8(device); + register_convrot(device); + register_rope(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/rope.cu b/src/core/cuda/rope.cu new file mode 100644 index 0000000000..1b724b0a78 --- /dev/null +++ b/src/core/cuda/rope.cu @@ -0,0 +1,337 @@ +#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 && 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; + } + 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 == 5, 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..4de48286ff --- /dev/null +++ b/src/core/cuda/sageattention/mma.cuh @@ -0,0 +1,683 @@ +/* + * 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 + +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..066f7fe082 --- /dev/null +++ b/src/core/cuda/sageattention/numeric_conversion.cuh @@ -0,0 +1,148 @@ +/* + * 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 + +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.h b/src/core/ggml_extend.h index 5747d625a6..4a97704d3e 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -4,6 +4,7 @@ #include #include +#include "core/ggml_extensions.h" #include "ggml-backend.h" #include "ggml.h" diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 8586d8766c..47fe80ffd8 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 @@ -17,9 +18,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) { @@ -450,24 +450,6 @@ bool sd_backend_is_cpu(ggml_backend_t backend) { return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU; } -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")) { @@ -596,7 +578,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; } bool sd_backend_cpu_set_n_threads(ggml_backend_t backend, int n_threads) { @@ -641,23 +625,6 @@ void sd_backend_clear_buffer_cache(ggml_backend_t backend) { } } -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; -} - ggml_status sd_backend_graph_compute_with_eval_callback(ggml_backend_t backend, ggml_cgraph* gf, sd_graph_eval_callback_t callback_eval, @@ -680,8 +647,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; } @@ -1207,6 +1173,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_extensions.cpp b/src/core/ggml_extensions.cpp new file mode 100644 index 0000000000..1f0aa1de50 --- /dev/null +++ b/src/core/ggml_extensions.cpp @@ -0,0 +1,364 @@ +#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 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..49b0e76059 --- /dev/null +++ b/src/core/ggml_extensions.h @@ -0,0 +1,66 @@ +#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_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 8454d43558..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,14 +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)); } - // Fusion checks need the original use counts, including consumers outside this segment. - for (ggml_tensor* tensor : represented) { - const size_t source_index = ggml_hash_find(&gf->visited_hash_set, tensor); - if (source_index != GGML_HASHSET_FULL && ggml_bitset_get(gf->visited_hash_set.used, source_index)) { - const size_t index = ggml_hash_insert(&segment_graph->visited_hash_set, tensor); - segment_graph->use_counts[index] = gf->use_counts[source_index]; - } - } + GGML_ASSERT(ggml_graph_cpy_use_counts(gf, segment_graph)); *graph_ctx_out = graph_ctx; return segment_graph; } @@ -896,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/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt new file mode 100644 index 0000000000..8e55d99970 --- /dev/null +++ b/src/core/vulkan/CMakeLists.txt @@ -0,0 +1,69 @@ +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(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") +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..4caadfbeff --- /dev/null +++ b/src/core/vulkan/int8.cpp @@ -0,0 +1,806 @@ +#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 = quant->src[0]->op == GGML_OP_ADD; + const int count = affine ? 4 : 3; + if (affine) { + const ggml_tensor* shift = quant->src[0]; + const ggml_tensor* add = shift->src[0]; + if (add->op != GGML_OP_ADD || add->src[1]->op != GGML_OP_MUL || add->src[1]->src[0] != add->src[0]) { + return 0; + } + nodes[0] = add->src[1]; + nodes[1] = add; + nodes[2] = shift; + nodes[3] = quant; + mod = nodes[0]->src[1]; + const ggml_tensor* vector = shift->src[1]; + if (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) || + vector->ne[0] != add->ne[0] || ggml_nelements(vector) != add->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) != 4) { + 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 ? 6 : 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[6]; + if (match_rms_modulate_convrot(device, graph, i + 5, indices) != 6) { + return 0; + } + for (int j = 0; j < 6; ++j) { + if (indices[j] != i + j) { + return 0; + } + } + return 6; + } + if (op != GGML_OP_NORM && op != GGML_OP_MUL) { + return 0; + } + const int count = op == GGML_OP_NORM ? 3 : 4; + 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 || first->op == GGML_OP_MUL || first->op == GGML_OP_RMS_NORM); + const bool rms = modulate && first->op == GGML_OP_RMS_NORM; + const bool affine = modulate && (rms || first->op == GGML_OP_MUL); + 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 + 2]->src[1] : 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}); + 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}, 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)}; + pipeline.constants = constants; + pipeline.constant_count = 2; + 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.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_MUL, "sd.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/rope.cpp b/src/core/vulkan/rope.cpp new file mode 100644 index 0000000000..63b976e79a --- /dev/null +++ b/src/core/vulkan/rope.cpp @@ -0,0 +1,226 @@ +#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 && 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->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 && copy->type == GGML_TYPE_F16 && + ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(scale, 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; + return 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)) ? 5 : 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* scale = output->op == GGML_OP_CPY ? output->src[0] : nullptr; + const auto* rope = scale && scale->op == GGML_OP_SCALE ? scale->src[0] : 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) ? 5 + : 0; + if (!count) + continue; + const ggml_tensor* nodes[] = {norm, mul, rope, scale, 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}; + 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 == 5 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan"; + pipeline.spirv = count == 5 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply; + pipeline.spirv_size = count == 5 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply); + 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_convrot.comp b/src/core/vulkan/shaders/modulate_convrot.comp new file mode 100644 index 0000000000..3e4b3a217e --- /dev/null +++ b/src/core/vulkan/shaders/modulate_convrot.comp @@ -0,0 +1,160 @@ +#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(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; + 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_rope_apply.comp b/src/core/vulkan/shaders/rms_rope_apply.comp new file mode 100644 index 0000000000..7d08f74149 --- /dev/null +++ b/src/core/vulkan/shaders/rms_rope_apply.comp @@ -0,0 +1,108 @@ +#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 +#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 + 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/model/common/rope.hpp b/src/model/common/rope.hpp index 63584bc5cc..b6fdd328e1 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -1077,8 +1077,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/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 739ef77f13..3dd720cad3 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -180,8 +180,8 @@ namespace MiniMaxH3 { 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 = ggml_rope_apply(ctx, x, pe, false); - if (ggml_backend_supports_op(runner_ctx->backend, out)) { + 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; } } 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 From 5356982584d5b46a8f5cb0842927aff6443fe174 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 1 Nov 2026 22:08:29 +0800 Subject: [PATCH 02/57] fix: chunk MiniMax-H3 MLP to respect backend buffer limits (#2257) --- src/model/diffusion/minimax_h3.hpp | 42 +++++++++++++++++++++++++----- 1 file changed, 36 insertions(+), 6 deletions(-) diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 3dd720cad3..0915b52d9c 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -140,8 +140,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); } @@ -159,11 +162,38 @@ namespace MiniMaxH3 { } 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 projected = fc1->forward(ctx, x); - auto activated = ggml_swiglu(ctx->ggml_ctx, projected); - return fc2->forward(ctx, activated); + 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)); + } + return ggml_ext_vec_concat(ctx->ggml_ctx, chunks, 1); } }; From df0fa26a61301b608c6baadd476f8337dc509832 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 1 Nov 2026 22:22:04 +0800 Subject: [PATCH 03/57] fix: update CUTLASS license paths in Docker and Windows CUDA packaging --- .github/workflows/build.yml | 2 +- docker/Dockerfile.cuda | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index ff7ca44984..c6c129ecf9 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -440,7 +440,7 @@ jobs: throw "Can't find stable-diffusion.dll" } if ("${{ matrix.build }}" -eq "cuda12") { - Copy-Item ggml/src/ggml-cuda/vendors/cutlass/LICENSE.txt (Join-Path (Split-Path -Parent $filePath) 'cutlass.txt') + Copy-Item thirdparty/cutlass/LICENSE.txt (Join-Path (Split-Path -Parent $filePath) 'cutlass.txt') } 7z a sd-${{ env.BRANCH_NAME }}-${{ steps.commit.outputs.short }}-bin-win-${{ matrix.build }}-x64.zip $filePath .\start-server.sh .\start-server.bat diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 13b4722f17..ef539a1a4b 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -53,7 +53,7 @@ RUN apt-get update && \ apt-get clean COPY --from=build /sd.cpp/build/bin /sd.cpp/bin -COPY --from=build /sd.cpp/ggml/src/ggml-cuda/vendors/cutlass/LICENSE.txt /sd.cpp/bin/cutlass.txt +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 From fe2682ba16dd3092054237e1c0ca1adf9107c73e Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 1 Nov 2026 22:55:56 +0800 Subject: [PATCH 04/57] fix: retry Vulkan capacity checks after memory reclamation (#2256) --- src/model_manager.cpp | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/src/model_manager.cpp b/src/model_manager.cpp index d07d63a810..8b6182797f 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1,9 +1,11 @@ #include "model_manager.h" #include +#include #include #include #include +#include #include #include @@ -2006,10 +2008,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)); From 62c54439c8bb4da32cab1d81981e7d447e2c93f3 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 2 Nov 2026 00:17:49 +0800 Subject: [PATCH 05/57] fix: include missing cassert headers in SageAttention CUDA code --- src/core/cuda/sageattention/mma.cuh | 1 + src/core/cuda/sageattention/numeric_conversion.cuh | 1 + 2 files changed, 2 insertions(+) diff --git a/src/core/cuda/sageattention/mma.cuh b/src/core/cuda/sageattention/mma.cuh index 4de48286ff..f473fa7296 100644 --- a/src/core/cuda/sageattention/mma.cuh +++ b/src/core/cuda/sageattention/mma.cuh @@ -22,6 +22,7 @@ #include #include #include +#include #include namespace sageattention { diff --git a/src/core/cuda/sageattention/numeric_conversion.cuh b/src/core/cuda/sageattention/numeric_conversion.cuh index 066f7fe082..8da67d5f13 100644 --- a/src/core/cuda/sageattention/numeric_conversion.cuh +++ b/src/core/cuda/sageattention/numeric_conversion.cuh @@ -22,6 +22,7 @@ #include #include #include +#include namespace sageattention { From 5d33b1bd2bada9c99f5e4c9d6765c915a2dcd452 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 2 Nov 2026 00:23:44 +0800 Subject: [PATCH 06/57] ci: trigger builds on CUDA header changes --- .github/workflows/build.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index c6c129ecf9..e310af4035 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -23,6 +23,7 @@ on: "**/*.c", "**/*.cpp", "**/*.cu", + "**/*.cuh", "ggml", "examples/server/frontend", "examples/server/frontend/**", @@ -44,6 +45,7 @@ on: "**/*.c", "**/*.cpp", "**/*.cu", + "**/*.cuh", "ggml", "examples/server/frontend", "examples/server/frontend/**", From 133087deb82abd8bcd4afdf0b895cf63a8714293 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 2 Nov 2026 01:09:36 +0800 Subject: [PATCH 07/57] fix: prefer discrete GPUs in auto-fit device selection (#2259) --- docs/backend.md | 5 ++++- src/core/backend_fit.cpp | 27 +++++++++++++++++++-------- 2 files changed, 23 insertions(+), 9 deletions(-) diff --git a/docs/backend.md b/docs/backend.md index 49ca0632b6..da90e48968 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -135,7 +135,10 @@ regardless of argument order, even with `--auto-fit on`. Auto-fit preserves explicit `--backend` assignments, including per-module assignments and device lists. For modules without a runtime assignment, it chooses -the GPU with the largest available memory budget (the first device on a tie). +the discrete GPU with the largest positive memory budget (the first device on a +tie), falling back to integrated GPUs when no discrete GPU has a positive budget. +An integrated GPU's shared system memory does not give it priority over a discrete +GPU. Explicit assignments such as `--backend Vulkan1` still select that device. It then derives parameter placements from the model metadata, each module's compute devices, and the remaining memory budgets. The chosen backend specifications are printed. diff --git a/src/core/backend_fit.cpp b/src/core/backend_fit.cpp index c970d10e68..120cbe423f 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 { @@ -140,6 +141,7 @@ namespace sd::backend_fit { Device device; device.name = ggml_backend_dev_name(dev); device.description = ggml_backend_dev_description(dev); + device.type = type; size_t free_bytes = 0, total_bytes = 0; ggml_backend_dev_memory(dev, &free_bytes, &total_bytes); device.free_bytes = (int64_t)free_bytes; @@ -193,14 +195,23 @@ namespace sd::backend_fit { } static size_t select_main_device(const std::vector& 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, From 290148d2e31ca6ef9757d567905776b5f3646db5 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 2 Nov 2026 01:19:28 +0800 Subject: [PATCH 08/57] feat: add custom modulate operator and ConvRot fusion (#2258) --- CMakeLists.txt | 2 +- cmake/backend_extensions.h.in | 1 + docs/ggml_extensions.md | 24 +++- src/core/cpu/modulate.cpp | 82 +++++++++++ src/core/cpu/registration.cpp | 1 + src/core/cpu/registration.h | 1 + src/core/cuda/CMakeLists.txt | 2 +- src/core/cuda/convrot.cu | 54 ++++--- src/core/cuda/convrot_kernel.h | 9 +- src/core/cuda/modulate.cu | 95 +++++++++++++ src/core/cuda/registration.cu | 2 + src/core/ggml_extend.cpp | 21 +++ src/core/ggml_extend.h | 9 ++ src/core/ggml_extensions.cpp | 77 ++++++++++ src/core/ggml_extensions.h | 4 + src/core/vulkan/CMakeLists.txt | 3 +- src/core/vulkan/int8.cpp | 56 ++++---- src/core/vulkan/modulate.cpp | 134 ++++++++++++++++++ src/core/vulkan/shaders/modulate.comp | 39 +++++ src/core/vulkan/shaders/modulate_convrot.comp | 5 +- src/model/diffusion/anima.hpp | 4 +- src/model/diffusion/boogu.hpp | 16 +-- src/model/diffusion/ernie_image.hpp | 6 +- src/model/diffusion/fibo2.hpp | 4 +- src/model/diffusion/flux.hpp | 22 ++- src/model/diffusion/flux3.hpp | 2 +- src/model/diffusion/ideogram4.hpp | 12 +- src/model/diffusion/krea2.hpp | 14 +- src/model/diffusion/lens.hpp | 10 +- src/model/diffusion/lingbot_video.hpp | 6 +- src/model/diffusion/ltxv.hpp | 38 ++--- src/model/diffusion/minimax_h3.hpp | 4 +- src/model/diffusion/mmdit.hpp | 24 ++-- src/model/diffusion/pixart.hpp | 6 +- src/model/diffusion/qwen_image.hpp | 10 +- src/model/diffusion/z_image.hpp | 17 ++- 36 files changed, 655 insertions(+), 161 deletions(-) create mode 100644 src/core/cpu/modulate.cpp create mode 100644 src/core/cuda/modulate.cu create mode 100644 src/core/vulkan/modulate.cpp create mode 100644 src/core/vulkan/shaders/modulate.comp diff --git a/CMakeLists.txt b/CMakeLists.txt index 235556facd..15bd91742e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -378,7 +378,7 @@ if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA)) 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) + list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8 sd_ggml_register_vulkan_modulate) endif() if(SD_BACKEND_REGISTRATIONS) string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}") diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in index cf41592fbe..6b42ea117f 100644 --- a/cmake/backend_extensions.h.in +++ b/cmake/backend_extensions.h.in @@ -5,6 +5,7 @@ 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_modulate(ggml_backend_dev_t device); static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = { @SD_BACKEND_REGISTRATIONS@ diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 6787c443ca..5c059892e8 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -135,6 +135,25 @@ CPU kernels and fusion rules live in `src/core/cpu/`, with registration collecte 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 @@ -269,8 +288,9 @@ 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`, `RMSNorm -> Mul(weight) -> Mul(modulation) -> Add -> -Add(shift) -> ConvRot`, `SwiGLU -> ConvRot`, and `SiLU/Sigmoid -> Mul -> 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. 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 index a4f5bd32be..d63fd1f099 100644 --- a/src/core/cpu/registration.cpp +++ b/src/core/cpu/registration.cpp @@ -6,6 +6,7 @@ namespace sd_cpu { void register_extensions(ggml_backend_dev_t device) { register_i8(device); register_rope(device); + register_modulate(device); } } diff --git a/src/core/cpu/registration.h b/src/core/cpu/registration.h index c711f6fac9..eba629c2a6 100644 --- a/src/core/cpu/registration.h +++ b/src/core/cpu/registration.h @@ -7,6 +7,7 @@ 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 diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 7bca6aa33f..f366899fe2 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW) function(sd_cuda_sources VENDOR LANGUAGE) set(SOURCES) - foreach(OP rope int8 convrot affine registration) + foreach(OP rope int8 convrot affine modulate 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") diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu index d5d159be14..cf7896114d 100644 --- a/src/core/cuda/convrot.cu +++ b/src/core/cuda/convrot.cu @@ -27,6 +27,7 @@ namespace SD_CUDA_NAMESPACE { } enum Kind { SCALE, NORM, + MOD, RMS_MOD, SWIGLU, SILU, @@ -65,16 +66,20 @@ namespace SD_CUDA_NAMESPACE { op.eps = param_f32(first, 0); if (op.k < 1024 || !vector(op.modulation, op.k)) return false; - } else if (count == 6 && first->op == GGML_OP_RMS_NORM && nodes[1]->op == GGML_OP_MUL && nodes[1]->src[0] == first && - nodes[2]->op == GGML_OP_MUL && nodes[2]->src[0] == nodes[1] && - nodes[3]->op == GGML_OP_ADD && nodes[3]->src[0] == nodes[1] && nodes[3]->src[1] == nodes[2] && - nodes[4]->op == GGML_OP_ADD && nodes[4]->src[0] == nodes[3] && value == nodes[4]) { + } 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[4]->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) || !vector(op.shift, op.k)) + 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; @@ -119,7 +124,7 @@ namespace SD_CUDA_NAMESPACE { 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 ? 6 : node->op == GGML_OP_NORM || node->op == GGML_OP_UNARY ? 3 + 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; @@ -132,7 +137,7 @@ namespace SD_CUDA_NAMESPACE { static void optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) { if (fusion_disabled()) return; - for (int i = 2; i < ggml_graph_n_nodes(graph); ++i) { + 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) @@ -140,17 +145,18 @@ namespace SD_CUDA_NAMESPACE { ggml_tensor* chain[6] = {}; int count = 0; auto* value = quant->src[0]; - if (value->op == GGML_OP_ADD && value->src[0]->op == GGML_OP_ADD) { - auto* add = value->src[0]; - auto* weight = add->src[0]; - if (weight->op == GGML_OP_MUL) { + 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] = add->src[1]; - chain[3] = add; - chain[4] = value; - chain[5] = quant; - count = 6; + 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) @@ -210,13 +216,13 @@ namespace SD_CUDA_NAMESPACE { 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 + template static CUDAError launch(Plan& p, const Operation& op, int8_t* output, CUDAStream stream) { - constexpr bool warp = Scaled || RMS; + constexpr bool warp = Scaled || RMS || Modulate; constexpr int extra = Normalize ? 384 : RMS ? 256 : 128; if (!p.configured) { - auto status = cuda_set_shared_memory(quantize, + auto status = cuda_set_shared_memory(quantize, int(p.device.shared_memory - extra)); if (status != cuda_success) return status; @@ -229,7 +235,7 @@ namespace SD_CUDA_NAMESPACE { : 256; const int64_t padded = (op.rows + 3) / 4 * 4; auto data = [](const ggml_tensor* t) { return t ? static_cast(t->data) : nullptr; }; - quantize<<>>( + 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(); @@ -260,6 +266,9 @@ namespace SD_CUDA_NAMESPACE { 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); @@ -277,12 +286,13 @@ namespace SD_CUDA_NAMESPACE { } 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}) { + 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; diff --git a/src/core/cuda/convrot_kernel.h b/src/core/cuda/convrot_kernel.h index 748c08e7de..04abf97ac3 100644 --- a/src/core/cuda/convrot_kernel.h +++ b/src/core/cuda/convrot_kernel.h @@ -40,7 +40,7 @@ namespace SD_CUDA_NAMESPACE { return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d); } - template + template __global__ void quantize( const T* src, int8_t* dst, @@ -103,7 +103,12 @@ namespace SD_CUDA_NAMESPACE { float value = src_row[i]; if constexpr (rms_modulation) { value = __fmul_rn(__fmul_rn(value, rms_scale), rms_weight[i]); - value = __fadd_rn(__fadd_rn(value, __fmul_rn(value, modulation[i])), rms_shift[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); 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 index d03b114cca..291edd62a6 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -11,6 +11,7 @@ namespace SD_CUDA_NAMESPACE { 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); #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); @@ -20,6 +21,7 @@ namespace SD_CUDA_NAMESPACE { if (cuda_device_info(device).device < 0) return; register_affine(device); + register_modulate(device); register_i8(device); register_convrot(device); register_rope(device); diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 6d72caedfb..ca1b231c86 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -6,6 +6,7 @@ #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 @@ -155,6 +156,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, diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 4a97704d3e..9ddfb0ad61 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -8,6 +8,8 @@ #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"); @@ -58,6 +60,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_extensions.cpp b/src/core/ggml_extensions.cpp index 1f0aa1de50..27089f6c3b 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -15,6 +15,83 @@ #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; diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index 49b0e76059..86c9121a7a 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -8,6 +8,10 @@ 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: diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt index 8e55d99970..42ee931a2e 100644 --- a/src/core/vulkan/CMakeLists.txt +++ b/src/core/vulkan/CMakeLists.txt @@ -48,6 +48,7 @@ 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_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) @@ -65,5 +66,5 @@ sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_C 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") +target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/rope.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/int8.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/modulate.cpp") target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}") diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index 4caadfbeff..7e3e67d48d 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -165,22 +165,15 @@ namespace sd_vulkan_i8 { const ggml_tensor* nodes[4]{}; const ggml_tensor* mod = nullptr; - const bool affine = quant->src[0]->op == GGML_OP_ADD; - const int count = affine ? 4 : 3; + const bool affine = sd_ggml_modulate_describe(quant->src[0]); + const int count = affine ? 2 : 3; if (affine) { - const ggml_tensor* shift = quant->src[0]; - const ggml_tensor* add = shift->src[0]; - if (add->op != GGML_OP_ADD || add->src[1]->op != GGML_OP_MUL || add->src[1]->src[0] != add->src[0]) { - return 0; - } - nodes[0] = add->src[1]; - nodes[1] = add; - nodes[2] = shift; - nodes[3] = quant; + nodes[0] = quant->src[0]; + nodes[1] = quant; mod = nodes[0]->src[1]; - const ggml_tensor* vector = shift->src[1]; - if (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) || - vector->ne[0] != add->ne[0] || ggml_nelements(vector) != add->ne[0]) { + 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 { @@ -215,7 +208,7 @@ namespace sd_vulkan_i8 { } 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) != 4) { + 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]; @@ -236,27 +229,27 @@ namespace sd_vulkan_i8 { } } - return next < 0 ? 6 : 0; + 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[6]; - if (match_rms_modulate_convrot(device, graph, i + 5, indices) != 6) { + int indices[4]; + if (match_rms_modulate_convrot(device, graph, i + 3, indices) != 4) { return 0; } - for (int j = 0; j < 6; ++j) { + for (int j = 0; j < 4; ++j) { if (indices[j] != i + j) { return 0; } } - return 6; + return 4; } - if (op != GGML_OP_NORM && op != GGML_OP_MUL) { + if (op != GGML_OP_NORM && !sd_ggml_modulate_describe(node(graph, i))) { return 0; } - const int count = op == GGML_OP_NORM ? 3 : 4; + 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; @@ -491,23 +484,24 @@ namespace sd_vulkan_i8 { 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 || first->op == GGML_OP_MUL || first->op == GGML_OP_RMS_NORM); + 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 || first->op == GGML_OP_MUL); + 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 + 2]->src[1] : mod) : 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}); - inputs.push_back({shift}); + if (shift) + inputs.push_back({shift}); if (rms) inputs.push_back({weight}); } @@ -521,7 +515,7 @@ namespace sd_vulkan_i8 { 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}, packed, scales}; + 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}; @@ -534,9 +528,9 @@ namespace sd_vulkan_i8 { *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)}; + const uint32_t constants[] = {k, uint32_t(affine), uint32_t(shift != nullptr)}; pipeline.constants = constants; - pipeline.constant_count = 2; + 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}}; @@ -788,6 +782,9 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) { 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; @@ -795,7 +792,6 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) { } fusions[] = { {GGML_OP_NORM, "sd.norm_convrot.vulkan.v1"}, {GGML_OP_RMS_NORM, "sd.rms_modulate_convrot.vulkan.v1"}, - {GGML_OP_MUL, "sd.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) { 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/shaders/modulate.comp b/src/core/vulkan/shaders/modulate.comp new file mode 100644 index 0000000000..df07626d2a --- /dev/null +++ b/src/core/vulkan/shaders/modulate.comp @@ -0,0 +1,39 @@ +#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 index 3e4b3a217e..3cac110c93 100644 --- a/src/core/vulkan/shaders/modulate_convrot.comp +++ b/src/core/vulkan/shaders/modulate_convrot.comp @@ -4,6 +4,7 @@ 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[]; @@ -111,7 +112,9 @@ void main() { if (AFFINE) { precise float modulated = values[col] * data_m[p.m_offset + col]; value = values[col] + modulated; - value += data_h[p.h_offset + col]; + if (HAS_SHIFT) { + value += data_h[p.h_offset + col]; + } } else { value = (values[col] - mean) * inv_std; value *= data_m[p.m_offset + col]; diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index 66e1a0ba35..e9e33b88a1 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -129,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}; } @@ -169,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; } }; diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index ac67ff9cd7..3f652e1589 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -128,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) { @@ -186,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}; } }; @@ -231,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; } @@ -334,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 { @@ -512,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)); diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 59b06cc2f1..15badb9834 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -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; } }; diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp index a642cb4dff..0c42235f8f 100644 --- a/src/model/diffusion/fibo2.hpp +++ b/src/model/diffusion/fibo2.hpp @@ -262,7 +262,7 @@ namespace Fibo2 { } return m; }; - auto h = ZImage::modulate(c, norm("attention_norm1", x), select(0, false)); + 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) { @@ -270,7 +270,7 @@ namespace Fibo2 { 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(c, norm("ffn_norm1", x), select(2, false)); + 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))); } diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index c9e3f5faab..40edca0ed9 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -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; diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp index 6e4cd64084..fac1b1a567 100644 --- a/src/model/diffusion/flux3.hpp +++ b/src/model/diffusion/flux3.hpp @@ -83,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]; } diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 5a5cb000cc..270344dc97 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -103,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) { @@ -285,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)); @@ -337,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; } diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index d21574b63b..fba817aa67 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -465,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); @@ -487,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); @@ -503,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], @@ -511,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], @@ -595,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], diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 62ad6d3732..bcfefd97bf 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -237,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); @@ -247,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])); @@ -272,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); } }; diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index 1ea7b1c0cf..d15d754080 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -485,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)); @@ -587,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); diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index 23c5f5f294..32e38012bd 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -38,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); } @@ -821,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)); @@ -831,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); @@ -847,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; @@ -1199,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]); @@ -1250,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, @@ -1268,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, @@ -1291,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]; @@ -1304,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]; @@ -1313,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])); @@ -1720,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); } diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 0915b52d9c..a7c2f1f930 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -445,9 +445,7 @@ namespace MiniMaxH3 { 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_add(ctx->ggml_ctx, - ggml_add(ctx->ggml_ctx, value, ggml_mul(ctx->ggml_ctx, value, row[scale])), - row[shift]); + return ggml_ext_modulate(ctx, value, row[shift], row[scale]); }; ggml_tensor* qkv = nullptr; for (size_t i = 0; i < segments.size(); ++i) { diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index cdf65f300f..613cfc11a7 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -372,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 { @@ -455,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}}; @@ -488,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}}; @@ -533,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; @@ -565,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; @@ -756,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/pixart.hpp b/src/model/diffusion/pixart.hpp index 45f9e43174..dc6e2c2e7d 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -198,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); @@ -288,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] diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 4c473da357..6435222f59 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -350,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); @@ -363,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); @@ -407,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; } diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index 229c6aaee0..1c66c96b7f 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -294,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 { @@ -361,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)); @@ -505,14 +504,14 @@ namespace ZImage { auto gate_mlp = mods[3]; auto residual = c; - c = modulate(ctx->ggml_ctx, attention_norm1->forward(ctx, c), scale_msa); + 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->ggml_ctx, ffn_norm1->forward(ctx, c), scale_mlp); + 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)); @@ -544,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; From 1350cf5097b644dbb3abcbc8cdf4521fdfbccef3 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 2 Nov 2026 23:51:33 +0800 Subject: [PATCH 09/57] ci: cache CUDA compilation for Windows and Docker builds --- .github/workflows/build.yml | 34 +++++++++++++++++++++++++++++- docker/Dockerfile.cuda | 17 +++++++++++++-- scripts/ci-cuda-cache.sh | 41 +++++++++++++++++++++++++++++++++++++ 3 files changed, 89 insertions(+), 3 deletions(-) create mode 100644 scripts/ci-cuda-cache.sh diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index e310af4035..7e3f6c17bf 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -16,6 +16,7 @@ on: ".github/workflows/**", ".dockerignore", "docker/Dockerfile*", + "scripts/ci-cuda-cache.sh", "**/CMakeLists.txt", "**/Makefile", "**/*.h", @@ -38,6 +39,7 @@ on: ".github/workflows/**", ".dockerignore", "docker/Dockerfile*", + "scripts/ci-cuda-cache.sh", "**/CMakeLists.txt", "**/Makefile", "**/*.h", @@ -253,6 +255,7 @@ jobs: run: echo "short=${GITHUB_SHA:0:7}" >> "$GITHUB_OUTPUT" - name: Set up Docker Buildx + id: buildx uses: docker/setup-buildx-action@v4 - name: Log in to the container registry @@ -275,6 +278,21 @@ jobs: # if set to "true" but frees about 6 GB tool-cache: false + - name: Cache CUDA compilation + if: matrix.variant == 'cuda' + uses: actions/cache@v5 + with: + path: ${{ runner.temp }}/cuda-ccache + key: cuda-ccache-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('docker/Dockerfile.cuda') }}-${{ github.run_id }}-${{ github.run_attempt }} + restore-keys: | + cuda-ccache-${{ runner.os }}-${{ runner.arch }}-${{ hashFiles('docker/Dockerfile.cuda') }}- + + - name: Restore CUDA compiler cache mount + if: matrix.variant == 'cuda' + env: + CUDA_CACHE_BUILDER: ${{ steps.buildx.outputs.name }} + run: bash scripts/ci-cuda-cache.sh restore + - name: Build and push Docker image id: build-push uses: docker/build-push-action@v7 @@ -288,6 +306,12 @@ jobs: annotations: ${{ steps.meta.outputs.annotations }} build-args: ${{ matrix.build-args }} + - name: Export CUDA compiler cache mount + if: matrix.variant == 'cuda' + env: + CUDA_CACHE_BUILDER: ${{ steps.buildx.outputs.name }} + run: bash scripts/ci-cuda-cache.sh save + macOS-latest-cmake: runs-on: macos-latest @@ -365,7 +389,7 @@ jobs: - build: "cpu" defines: "-DGGML_NATIVE=OFF -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "cuda12" - defines: "-DSD_CUDA=ON -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" + defines: "-DSD_CUDA=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache -DCMAKE_CUDA_COMPILER_LAUNCHER=sccache -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "vulkan" defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" steps: @@ -394,6 +418,14 @@ jobs: method: "network" sub-packages: '["nvcc", "cudart", "cublas", "cublas_dev", "thrust", "visual_studio_integration"]' + - name: Cache CUDA compilation + if: matrix.build == 'cuda12' + uses: ggml-org/ccache-action@v1.2.21 + with: + variant: sccache + key: windows-cuda-12.8.1-x64 + max-size: "2G" + - name: Install Vulkan SDK id: get_vulkan if: ${{ matrix.build == 'vulkan' }} diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index ef539a1a4b..e3364e8402 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,3 +1,5 @@ +# syntax=docker/dockerfile:1 + ARG CUDA_VERSION=12.6.3 ARG UBUNTU_VERSION=24.04 @@ -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,7 +53,11 @@ 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}-runtime-ubuntu${UBUNTU_VERSION} AS runtime 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 From 250ea6e3857133d8d99c551490a31cdb37182b1e Mon Sep 17 00:00:00 2001 From: cumal <43264640+cumal@users.noreply.github.com> Date: Mon, 2 Nov 2026 17:03:43 +0100 Subject: [PATCH 10/57] fix: layer split across cuda (#2262) --- src/core/layer_split_partition.cpp | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index e19f47b261..6ef24357a4 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -206,23 +206,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); } } From b4d6b83af6d12e5770306f82b1e5625ed20ec237 Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 3 Nov 2026 00:07:40 +0800 Subject: [PATCH 11/57] fix: support integer LoRA scalar conversion to FP32 (#2263) --- src/model_loader.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/model_loader.cpp b/src/model_loader.cpp index c3682a77c3..5576067409 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) { From 3be23788fc0c9f35a2c24b1211caae11d4a54dbc Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 3 Nov 2026 00:20:09 +0800 Subject: [PATCH 12/57] ci: use ccache to fix Windows CUDA command parsing failures --- .github/workflows/build.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 7e3f6c17bf..7c407139bb 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -389,7 +389,7 @@ jobs: - build: "cpu" defines: "-DGGML_NATIVE=OFF -DSD_BUILD_SHARED_LIBS=ON -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "cuda12" - defines: "-DSD_CUDA=ON -DCMAKE_C_COMPILER_LAUNCHER=sccache -DCMAKE_CXX_COMPILER_LAUNCHER=sccache -DCMAKE_CUDA_COMPILER_LAUNCHER=sccache -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" + defines: "-DSD_CUDA=ON -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache -DSD_BUILD_SHARED_LIBS=ON -DCMAKE_CUDA_ARCHITECTURES='61;70;75;80;86;89;90;100;120' -DCMAKE_CUDA_FLAGS='-Xcudafe \"--diag_suppress=177\" -Xcudafe \"--diag_suppress=550\"' -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" - build: "vulkan" defines: "-DSD_VULKAN=ON -DSD_BUILD_SHARED_LIBS=ON -DGGML_NATIVE=OFF -DSD_BUILD_SHARED_GGML_LIB=ON -DGGML_BACKEND_DL=ON -DGGML_CPU_ALL_VARIANTS=ON" steps: @@ -422,7 +422,8 @@ jobs: if: matrix.build == 'cuda12' uses: ggml-org/ccache-action@v1.2.21 with: - variant: sccache + variant: ccache + install: binary key: windows-cuda-12.8.1-x64 max-size: "2G" From f8ea2e162c1bea4530aa775dd25009dfe41fddbe Mon Sep 17 00:00:00 2001 From: Mirek Date: Tue, 3 Nov 2026 15:41:09 +0000 Subject: [PATCH 13/57] perf: use native SiTU for MOSS with composite fallback (#2210) --- ggml | 2 +- src/core/ggml_extend.cpp | 21 +++++++++++++++++++++ src/core/ggml_extend.h | 3 +++ src/model/diffusion/moss_image.hpp | 12 +++--------- 4 files changed, 28 insertions(+), 10 deletions(-) diff --git a/ggml b/ggml index d49c49e1dd..ccb88f810b 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit d49c49e1dd361ba3c0a77c90b5b1404b83969aa8 +Subproject commit ccb88f810be05bccc9621e706354611ccc2462c7 diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index ca1b231c86..34d35056a9 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -144,6 +144,27 @@ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_firs return gate_first ? ggml_swiglu(ctx, x) : ggml_swiglu_swapped(ctx, 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, ggml_tensor* a) { const float eps = 1e-6f; // default eps parameter diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 9ddfb0ad61..791155780f 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -52,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); diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 188c74196d..36d8d8328b 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -4,6 +4,7 @@ #include #include #include + #include "core/pe_cache.h" #include "model/common/ggml_block.hpp" @@ -55,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; @@ -94,7 +88,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); } @@ -137,7 +131,7 @@ namespace MOSSImage { y = norm("norm2", x); y = ggml_add(c, ggml_add(c, y, ggml_mul(c, y, part(4))), part(3)); 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); From 6225bf8c88a76e4baf3b7be86c0e6933414f26eb Mon Sep 17 00:00:00 2001 From: "Piotr Wilkin (ilintar)" Date: Tue, 3 Nov 2026 16:46:35 +0100 Subject: [PATCH 14/57] perf: view qkv slices for moss instead of slice+reshape (#2232) --- src/model/diffusion/moss_image.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 36d8d8328b..1f5dc16685 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -79,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)); From d30854463dbc75fb3671209b79016bb0d8ac321e Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 3 Nov 2026 23:59:21 +0800 Subject: [PATCH 15/57] perf: fuse Z-Image GPU operations and optimize BF16 GEMM (#2265) --- CMakeLists.txt | 8 +- src/core/cuda/CMakeLists.txt | 3 +- src/core/cuda/bf16.cu | 277 ++++++++++++++++++++++++++++++++++ src/core/cuda/bf16.h | 26 ++++ src/core/cuda/bf16_gemm.cu | 25 +++ src/core/cuda/bf16_gemm.h | 13 ++ src/core/cuda/registration.cu | 4 + src/core/cuda/rms_affine.cu | 168 +++++++++++++++++++++ src/core/cuda/rms_affine.h | 11 ++ 9 files changed, 530 insertions(+), 5 deletions(-) create mode 100644 src/core/cuda/bf16.cu create mode 100644 src/core/cuda/bf16.h create mode 100644 src/core/cuda/bf16_gemm.cu create mode 100644 src/core/cuda/bf16_gemm.h create mode 100644 src/core/cuda/rms_affine.cu create mode 100644 src/core/cuda/rms_affine.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 15bd91742e..048f273dee 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -350,15 +350,15 @@ if (SD_CUDA) 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) - 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 PROPERTIES + 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 APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + 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 APPEND PROPERTY COMPILE_DEFINITIONS SD_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() diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index f366899fe2..0a95d6cf7c 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW) function(sd_cuda_sources VENDOR LANGUAGE) set(SOURCES) - foreach(OP rope int8 convrot affine modulate registration) + 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") @@ -27,6 +27,7 @@ if(SD_CUDA) 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") 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/registration.cu b/src/core/cuda/registration.cu index 291edd62a6..d75728021f 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -12,6 +12,8 @@ namespace SD_CUDA_NAMESPACE { 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); @@ -25,6 +27,8 @@ namespace SD_CUDA_NAMESPACE { 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); diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu new file mode 100644 index 0000000000..d9ac351edf --- /dev/null +++ b/src/core/cuda/rms_affine.cu @@ -0,0 +1,168 @@ +#include +#include +#include +#include "bf16.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 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 + 1 < ggml_graph_n_nodes(graph); ++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* affine = mul->src[0]; + if (affine->op != GGML_OP_MUL || 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; + } + + template + static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) { + const int col0 = threadIdx.x; + const int64_t base = int64_t(blockIdx.x) * width; + float values[8], sum = 0.f; + for (int j = 0, col = col0; col < width; ++j, col += Block) { + values[j] = x[base + col]; + sum = __fmaf_rn(values[j], values[j], sum); + } + sum = warp_sum(sum); + __shared__ float sums[32]; + const int lane = col0 % 32; + if (lane == 0) + sums[col0 / 32] = sum; + __syncthreads(); + sum = warp_sum(lane < Block / 32 ? sums[lane] : 0.f); + const float scale = rsqrtf(sum / width + eps); + for (int j = 0, col = col0; col < width; ++j, col += Block) { + const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]); + const float r = Modulate ? value : residual[base + col]; + dst[base + col] = T(__fadd_rn(__fmul_rn(value, gate[col]), r)); + } + } + + 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, void**, size_t* workspace, void*) { + *workspace = ggml_nbytes(nodes[3]); + return GGML_STATUS_SUCCESS; + } + static void destroy(void*, void*) {} + + static ggml_status 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; + const auto* x = nodes[0]->src[0]; + const auto* w = nodes[1]->src[1]; + const auto* gate = nodes[2]->src[1]; + auto* dst = nodes[3]; + const auto* residual = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0]; + const bool modulation = residual == nodes[1]; + const bool alias = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) || + (!modulation && overlaps(dst, residual) && dst->data != residual->data); + auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); + float eps; + std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); + auto launch = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>) + : (modulation ? kernel<1024, true> : kernel<1024, false>); + const auto stream = static_cast(context.stream); + const int threads = x->ne[0] < 1024 ? 256 : 1024; + launch<<>>( + static_cast(x->data), static_cast(w->data), static_cast(gate->data), + modulation ? nullptr : static_cast(residual->data), output, int(x->ne[0]), eps); + 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; + } + } + + 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* x = nodes[0]->src[0]; + const auto* residual = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0]; + const bool modulation = residual == nodes[1]; + auto launch = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>) + : (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>); + float eps; + std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); + const int threads = x->ne[0] < 1024 ? 256 : 1024; + launch<<>>( + static_cast(x->data), static_cast(nodes[1]->src[1]->data), + static_cast(nodes[2]->src[1]->data), modulation ? nullptr : static_cast(residual->data), + static_cast(output), int(x->ne[0]), eps); + return cuda_last_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)); + } +} 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 From 984eb0686cbf5b7fc4f0063d9e0b1fb38d5c4f82 Mon Sep 17 00:00:00 2001 From: leejet Date: Wed, 4 Nov 2026 02:01:41 +0800 Subject: [PATCH 16/57] perf(cuda): fuse MOSS modulation and direct RoPE F16 output (#2269) Co-authored-by: mirek190 --- src/core/cuda/rms_affine.cu | 206 +++++++++++++++++++++-------- src/core/cuda/rope.cu | 11 +- src/model/diffusion/moss_image.hpp | 6 +- 3 files changed, 164 insertions(+), 59 deletions(-) diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu index d9ac351edf..9f4192b259 100644 --- a/src/core/cuda/rms_affine.cu +++ b/src/core/cuda/rms_affine.cu @@ -2,6 +2,7 @@ #include #include #include "bf16.h" +#include "core/ggml_extensions.h" #include "rms_affine.h" namespace SD_CUDA_NAMESPACE { @@ -18,6 +19,28 @@ namespace SD_CUDA_NAMESPACE { 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)) @@ -44,13 +67,17 @@ namespace SD_CUDA_NAMESPACE { static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) { if (disabled()) return; - for (int i = 2; i + 1 < ggml_graph_n_nodes(graph); ++i) { + for (int i = 2; i < ggml_graph_n_nodes(graph); ++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; + 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->op != GGML_OP_MUL || affine->src[0]->op != GGML_OP_RMS_NORM) + 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) { @@ -68,64 +95,136 @@ namespace SD_CUDA_NAMESPACE { return value; } - template - static __global__ void kernel(const float* x, const float* weight, const float* gate, const float* residual, T* dst, int width, float eps) { + 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[8], sum = 0.f; - for (int j = 0, col = col0; col < width; ++j, col += Block) { - values[j] = x[base + col]; - sum = __fmaf_rn(values[j], values[j], sum); - } - sum = warp_sum(sum); + float values[Groups][8]; __shared__ float sums[32]; const int lane = col0 % 32; - if (lane == 0) - sums[col0 / 32] = sum; +#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(); - sum = warp_sum(lane < Block / 32 ? sums[lane] : 0.f); + const float sum = warp_sum(lane < Block * Groups / 32 ? sums[lane] : 0.f); const float scale = rsqrtf(sum / width + eps); - for (int j = 0, col = col0; col < width; ++j, col += Block) { - const float value = __fmul_rn(__fmul_rn(scale, values[j]), weight[col]); - const float r = Modulate ? value : residual[base + col]; - dst[base + col] = T(__fadd_rn(__fmul_rn(value, gate[col]), r)); +#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, void**, size_t* workspace, void*) { - *workspace = ggml_nbytes(nodes[3]); + 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, 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* x = nodes[0]->src[0]; - const auto* w = nodes[1]->src[1]; - const auto* gate = nodes[2]->src[1]; - auto* dst = nodes[3]; - const auto* residual = dst->src[0] == nodes[2] ? dst->src[1] : dst->src[0]; - const bool modulation = residual == nodes[1]; - const bool alias = (overlaps(dst, x) && dst->data != x->data) || overlaps(dst, w) || overlaps(dst, gate) || - (!modulation && overlaps(dst, residual) && dst->data != residual->data); - auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); - float eps; - std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); - auto launch = x->ne[0] < 1024 ? (modulation ? kernel<256, true> : kernel<256, false>) - : (modulation ? kernel<1024, true> : kernel<1024, false>); + 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); - const int threads = x->ne[0] < 1024 ? 256 : 1024; - launch<<>>( - static_cast(x->data), static_cast(w->data), static_cast(gate->data), - modulation ? nullptr : static_cast(residual->data), output, int(x->ne[0]), eps); - auto error = cuda_last_error(); + 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; @@ -137,19 +236,12 @@ namespace SD_CUDA_NAMESPACE { } ggml_status rms_affine_pack(CUDAStream stream, ggml_tensor* const* nodes, void* output) { - const auto* x = nodes[0]->src[0]; - const auto* residual = nodes[3]->src[0] == nodes[2] ? nodes[3]->src[1] : nodes[3]->src[0]; - const bool modulation = residual == nodes[1]; - auto launch = x->ne[0] < 1024 ? (modulation ? sd_rms_affine::kernel<256, true, CUDABFloat16> : sd_rms_affine::kernel<256, false, CUDABFloat16>) - : (modulation ? sd_rms_affine::kernel<1024, true, CUDABFloat16> : sd_rms_affine::kernel<1024, false, CUDABFloat16>); - float eps; - std::memcpy(&eps, nodes[0]->op_params, sizeof(eps)); - const int threads = x->ne[0] < 1024 ? 256 : 1024; - launch<<>>( - static_cast(x->data), static_cast(nodes[1]->src[1]->data), - static_cast(nodes[2]->src[1]->data), modulation ? nullptr : static_cast(residual->data), - static_cast(output), int(x->ne[0]), eps); - return cuda_last_error() == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + 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) { @@ -164,5 +256,11 @@ namespace SD_CUDA_NAMESPACE { 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/rope.cu b/src/core/cuda/rope.cu index 1b724b0a78..99cc824f1f 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -169,7 +169,7 @@ namespace SD_CUDA_NAMESPACE { return false; const auto* x = norm->src[0]; const auto* weight = mul->src[1]; - return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 && + 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] && @@ -232,6 +232,13 @@ namespace SD_CUDA_NAMESPACE { 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; } @@ -293,7 +300,7 @@ namespace SD_CUDA_NAMESPACE { 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 == 5, count == 5 ? param_f32(nodes[3], 0) : 1.f); + 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]; diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 1f5dc16685..a0fae7448f 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -116,7 +116,7 @@ namespace MOSSImage { auto norm = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; 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_view_4d(c, qkv, config.head_dim, config.dim / config.head_dim, x->ne[1], x->ne[2], @@ -130,7 +130,7 @@ 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 = ggml_ext_situ(ctx, up_gate); y = linear("proj.ffn_down", y); @@ -191,7 +191,7 @@ 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); } From 3e92bf990ad70566212ee5eb27ad7570fb7ed2af Mon Sep 17 00:00:00 2001 From: leejet Date: Wed, 4 Nov 2026 23:24:51 +0800 Subject: [PATCH 17/57] feat: add P-H3-Image generation and editing support (#2272) --- README.md | 1 + docs/p_h3_image.md | 34 ++++++++ examples/common/common.cpp | 8 +- examples/server/recipes/defaults.md | 1 + examples/server/recipes/p-h3-image.json | 55 +++++++++++++ src/conditioning/conditioner.hpp | 64 ++++++++++++--- src/core/layer_split_partition.cpp | 2 +- src/core/vulkan/shaders/modulate.comp | 19 +++-- src/core/vulkan/shaders/modulate_convrot.comp | 4 +- src/model.h | 3 +- src/model/integrations/minimax_h3.cpp | 9 +++ src/model/te/ph3_image_connector.hpp | 70 ++++++++++++++++ src/model/vae/minimax_h3_vae.hpp | 79 ++++++++++++++++--- src/model_loader.cpp | 19 +++++ src/name_conversion.cpp | 7 ++ src/pipeline/diffusion_engine.cpp | 18 +++-- src/pipeline/image.cpp | 57 ++++++++++++- src/pipeline/model_builders.cpp | 12 ++- 18 files changed, 421 insertions(+), 41 deletions(-) create mode 100644 docs/p_h3_image.md create mode 100644 examples/server/recipes/p-h3-image.json create mode 100644 src/model/te/ph3_image_connector.hpp diff --git a/README.md b/README.md index 1337c5b92f..28af00a697 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,7 @@ API and command-line option may change frequently.*** - [Z-Image](./docs/z_image.md) - [Z-Image L2P](./docs/z_image_l2p.md) - [MiniT2I](./docs/minit2i.md) + - [P-H3-Image](./docs/p_h3_image.md) - [SenseNova U1.5](./docs/sensenova_u1.md) - [Ovis-Image](./docs/ovis_image.md) - [Anima](./docs/anima.md) 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/examples/common/common.cpp b/examples/common/common.cpp index 637bbb659f..d5496f66fe 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}, {"", @@ -1530,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); @@ -1777,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", diff --git a/examples/server/recipes/defaults.md b/examples/server/recipes/defaults.md index e20a2cd783..efe596c92b 100644 --- a/examples/server/recipes/defaults.md +++ b/examples/server/recipes/defaults.md @@ -71,6 +71,7 @@ For community edit adapters and other optional extensions, these are the base-mo | `minit2i` | 2.5 | 20 | 512 Γ— 512 | [Source](https://huggingface.co/MiniT2I/MiniT2I/raw/main/pipeline.py) B/16 example and native model size; L/16 uses CFG 6. | | `moss-image-base` | 3 | 20 | 1024 Γ— 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-Base) | | `moss-image-sft` | 3 | 20 | 1024 Γ— 1024 | [Source](https://huggingface.co/OpenMOSS-Team/MOSS-Image-SFT) | +| `p-h3-image` | 1 | 8 | 1024 Γ— 1024 | [Runtime guide](../../../docs/p_h3_image.md) Fixed P-H3 sigma schedule; upstream 2912 Γ— 2912 reduced to the catalog default. | | `nitrosd-realism` | 1 | 1 | 1024 Γ— 1024 | [Source](https://huggingface.co/ChenDY/NitroFusion) | | `ovis-image` | 5 | 20 | 1024 Γ— 1024 | [Source](https://huggingface.co/ATH-MaaS/Ovis-Image-7B) | | `pid` | 1 | 4 | 2048 Γ— 2048 | [Source](https://huggingface.co/nvidia/PiD) Standalone PiD decoding at 2K; guidance is handled by the upstream generator. | 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/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index e3832a1d8e..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" @@ -1976,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 = {}, @@ -2022,12 +2024,18 @@ struct LLMEmbedder : public Conditioner { } else if (version == VERSION_MOSS_IMAGE) { arch = LLM::LLMArch::QWEN3_5; } - llm = std::make_shared(arch, - backend, - tensor_storage_map, - "text_encoders.llm", - enable_vision, - weight_manager); + 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; @@ -2063,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"); } @@ -2074,6 +2085,9 @@ struct LLMEmbedder : public Conditioner { void set_max_graph_vram_bytes(size_t max_vram_bytes) override { llm->set_max_graph_vram_bytes(max_vram_bytes); + if (ph3_connector) { + ph3_connector->set_max_graph_vram_bytes(max_vram_bytes); + } if (byt5) { byt5->set_max_graph_vram_bytes(max_vram_bytes); } @@ -2081,6 +2095,9 @@ struct LLMEmbedder : public Conditioner { void set_runtime_backends(const std::vector& backends) override { llm->set_runtime_backends(backends); + if (ph3_connector) { + ph3_connector->set_runtime_backends(backends); + } if (byt5) { byt5->set_runtime_backends(backends); } @@ -2113,6 +2130,9 @@ struct LLMEmbedder : public Conditioner { void set_flash_attention_enabled(bool enabled) override { llm->set_flash_attention_enabled(enabled); + if (ph3_connector) { + ph3_connector->set_flash_attention_enabled(enabled); + } if (byt5) { byt5->set_flash_attention_enabled(enabled); } @@ -2120,12 +2140,18 @@ struct LLMEmbedder : public Conditioner { void set_scale_overrides(float linear_scale, float attn_scale) override { llm->set_scale_overrides(linear_scale, attn_scale); + if (ph3_connector) { + ph3_connector->set_scale_overrides(linear_scale, attn_scale); + } if (byt5) { byt5->set_scale_overrides(linear_scale, attn_scale); } } void set_weight_adapter(const std::shared_ptr& adapter) override { + if (ph3_connector) { + ph3_connector->set_weight_adapter(adapter); + } if (llm) { llm->set_weight_adapter(adapter); } @@ -2135,6 +2161,9 @@ struct LLMEmbedder : public Conditioner { } void runner_end() override { + if (ph3_connector) { + ph3_connector->runner_end(); + } if (llm) { llm->runner_end(); } @@ -2311,6 +2340,12 @@ struct LLMEmbedder : public Conditioner { LOG_ERROR("LLM prompt encoding failed"); return {}; } + if (ph3_connector) { + hidden_states = ph3_connector->compute(n_threads, hidden_states); + if (hidden_states.empty()) { + return {}; + } + } hidden_states = apply_token_weights(std::move(hidden_states), weights); GGML_ASSERT(hidden_states.shape()[1] > prompt_template_encode_start_idx); @@ -2457,8 +2492,9 @@ struct LLMEmbedder : public Conditioner { if (sd_version_is_minimax_h3(version)) { prompt_template_encode_start_idx = 0; - out_layers = {50}; - prompt_attn_range = {0, 0}; + // The final tap is normalized; LLMRunner exposes it at num_layers + 1. + out_layers = version == VERSION_P_H3_IMAGE ? std::set{13, 19, 25, 31, 37} : std::set{50}; + prompt_attn_range = {0, 0}; if (llm->enable_vision) { const int patch_size = llm->config.vision.patch_size; @@ -2564,7 +2600,17 @@ struct LLMEmbedder : public Conditioner { } } } - prompt += conditioner_params.text; + auto text = conditioner_params.text; + if (version == VERSION_P_H3_IMAGE && !image_embeds.empty()) { + // Preserve explicit picture targets and structured briefs from the upstream upsampler. + static const std::regex verbatim(R"(^\s*(Edit |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}; diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index 6ef24357a4..ff8aaeb602 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -206,7 +206,7 @@ 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]; - auto assign_node = [&](int node_index) { + auto assign_node = [&](int node_index) { if (node_index < 0 || node_index >= n_nodes) { return; } diff --git a/src/core/vulkan/shaders/modulate.comp b/src/core/vulkan/shaders/modulate.comp index df07626d2a..81f910d07b 100644 --- a/src/core/vulkan/shaders/modulate.comp +++ b/src/core/vulkan/shaders/modulate.comp @@ -1,16 +1,25 @@ #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(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; +} +p; void main() { uint group = gl_WorkGroupID.y * gl_NumWorkGroups.x + gl_WorkGroupID.x; diff --git a/src/core/vulkan/shaders/modulate_convrot.comp b/src/core/vulkan/shaders/modulate_convrot.comp index 3cac110c93..4e85b1d0b8 100644 --- a/src/core/vulkan/shaders/modulate_convrot.comp +++ b/src/core/vulkan/shaders/modulate_convrot.comp @@ -2,8 +2,8 @@ #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 = 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 { diff --git a/src/model.h b/src/model.h index 12bc22e0f0..4695f5c74d 100644 --- a/src/model.h +++ b/src/model.h @@ -70,6 +70,7 @@ enum SDVersion { VERSION_MOSS_IMAGE, VERSION_FIBO2, VERSION_CHROMA, + VERSION_P_H3_IMAGE, VERSION_COUNT, }; @@ -166,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) { diff --git a/src/model/integrations/minimax_h3.cpp b/src/model/integrations/minimax_h3.cpp index 2569b626ea..35a1f6700b 100644 --- a/src/model/integrations/minimax_h3.cpp +++ b/src/model/integrations/minimax_h3.cpp @@ -40,6 +40,15 @@ const std::vector& minimax_h3_definitions() { 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/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/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 1e06558a19..d5bb01a4db 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -10,6 +10,7 @@ #include #include +#include "core/rng_mt19937.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" #include "model/vae/vae.hpp" @@ -18,6 +19,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; @@ -361,13 +390,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)] = @@ -389,7 +419,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, @@ -443,7 +475,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, @@ -451,7 +486,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, @@ -464,7 +499,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, @@ -483,6 +518,7 @@ namespace MiniMaxH3VAE { }; struct MiniMaxH3VideoVAERunner : public VAE { + MiniMaxH3VAEConfig config; MiniMaxH3VideoVAE model; sd::Tensor pixel_mean; sd::Tensor pixel_std; @@ -495,6 +531,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}, @@ -516,12 +554,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 { @@ -531,6 +569,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; } @@ -610,6 +657,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) { @@ -654,6 +705,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, @@ -752,7 +807,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_loader.cpp b/src/model_loader.cpp index 5576067409..407885bbff 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -578,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" && diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 57a41c5c39..2e147348ff 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1699,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."}, diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index d161d9e5be..98e80e1174 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -778,7 +778,7 @@ 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 (!load_model("embeddings_connectors", 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); } } @@ -1491,7 +1491,7 @@ bool StableDiffusionGGML::build_denoiser() { } case H3_AV_FLOW_PRED: { LOG_INFO("running in MiniMax H3 AV FLOW mode"); - denoiser = std::make_shared(default_flow_shift, 3.f, get_latent_channel()); + 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: { @@ -2594,9 +2594,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr 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(); @@ -3021,6 +3024,9 @@ sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensordecode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y); + 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(), @@ -3220,7 +3226,7 @@ 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 (sd_version_is_minimax_h3(version)) { + 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 diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 607ec18c5e..0763e38ca0 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -144,6 +144,34 @@ namespace sd::pipeline { 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, @@ -323,7 +351,10 @@ namespace sd::pipeline { 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; @@ -450,6 +481,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); @@ -502,6 +543,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() && @@ -546,6 +590,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; @@ -1047,13 +1094,19 @@ namespace sd::pipeline { 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); diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index 8b85f9b46b..e78a0d666e 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -110,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; @@ -159,6 +159,13 @@ namespace sd::model_builders { version, weight_manager); } 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", @@ -240,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) { From d1358945dc997bc97c702a3c90f59fa6bce5377d Mon Sep 17 00:00:00 2001 From: farawayso Date: Thu, 5 Nov 2026 02:59:04 +0800 Subject: [PATCH 18/57] feat: MiniMax-H3 latent hires second pass (#2249) --- docs/minimax_h3.md | 51 ++ src/model/upscaler/h3_latent_upscaler.hpp | 584 ++++++++++++++++++++++ src/pipeline/video.cpp | 316 +++++++++++- 3 files changed, 930 insertions(+), 21 deletions(-) create mode 100644 src/model/upscaler/h3_latent_upscaler.hpp diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index 0c17ba0e01..db554ddb03 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -99,6 +99,57 @@ 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/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/pipeline/video.cpp b/src/pipeline/video.cpp index a07b93175c..8c911b143a 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -12,6 +12,7 @@ #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" @@ -1317,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, @@ -1630,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"); - return false; - } - if (request.hires.upscaler != SD_HIRES_UPSCALER_MODEL) { - LOG_ERROR("LTX latent spatial upscale currently requires hires upscaler MODEL"); + 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 (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; + } } } @@ -1668,7 +1917,7 @@ namespace sd::pipeline { } 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); @@ -1780,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); @@ -1798,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); @@ -1829,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, @@ -1843,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, @@ -1852,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; @@ -1872,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], From 58bb3c1b1cd4749e94e39f537e9adfce295dcaee Mon Sep 17 00:00:00 2001 From: leejet Date: Thu, 5 Nov 2026 03:01:31 +0800 Subject: [PATCH 19/57] perf(vulkan): accelerate MOSS normalization and SiTU staging (#2273) Co-authored-by: mirek190 --- CMakeLists.txt | 2 +- cmake/backend_extensions.h.in | 1 + ggml | 2 +- src/core/vulkan/CMakeLists.txt | 3 +- src/core/vulkan/rms_affine.cpp | 192 ++++++++++++++++++++ src/core/vulkan/rope.cpp | 32 ++-- src/core/vulkan/shaders/rms_modulate.comp | 136 ++++++++++++++ src/core/vulkan/shaders/rms_rope_apply.comp | 7 +- 8 files changed, 358 insertions(+), 17 deletions(-) create mode 100644 src/core/vulkan/rms_affine.cpp create mode 100644 src/core/vulkan/shaders/rms_modulate.comp diff --git a/CMakeLists.txt b/CMakeLists.txt index 048f273dee..8c5866d494 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -378,7 +378,7 @@ if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA)) 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) + 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}") diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in index 6b42ea117f..2dc1e0f6cd 100644 --- a/cmake/backend_extensions.h.in +++ b/cmake/backend_extensions.h.in @@ -5,6 +5,7 @@ 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) = { diff --git a/ggml b/ggml index ccb88f810b..135ce1e0e3 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit ccb88f810be05bccc9621e706354611ccc2462c7 +Subproject commit 135ce1e0e3f2b0693d3de93d4f6b5fa0a1947820 diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt index 42ee931a2e..55acf0bb91 100644 --- a/src/core/vulkan/CMakeLists.txt +++ b/src/core/vulkan/CMakeLists.txt @@ -49,6 +49,7 @@ endfunction() # 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) @@ -66,5 +67,5 @@ sd_vulkan_shader(mul_mat_i8_swiglu_split_cm2 mul_mat_i8_tensorwise_cm2 ${SD_VK_C 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") +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/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 index 63b976e79a..194603e13d 100644 --- a/src/core/vulkan/rope.cpp +++ b/src/core/vulkan/rope.cpp @@ -72,7 +72,7 @@ namespace sd_vulkan_rope { return false; const auto* x = norm->src[0]; const auto* weight = mul->src[1]; - return x && weight && x->type == GGML_TYPE_F32 && x->ne[0] == 128 && x->nb[0] == 4 && + 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]; @@ -81,32 +81,35 @@ namespace sd_vulkan_rope { 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->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 && copy->type == GGML_TYPE_F16 && - ggml_is_contiguous(copy) && aligned(copy, 4) && ggml_are_same_shape(scale, copy); + 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; - return 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)) ? 5 : 3; + 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* scale = output->op == GGML_OP_CPY ? output->src[0] : nullptr; - const auto* rope = scale && scale->op == GGML_OP_SCALE ? scale->src[0] : output; + 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) ? 5 + 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, output}; + 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]) @@ -170,6 +173,7 @@ namespace sd_vulkan_rope { 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; @@ -181,9 +185,13 @@ namespace sd_vulkan_rope { groups[0] = std::min(n_groups, info.max_group_count[0]); groups[1] = (n_groups + groups[0] - 1) / groups[0]; } else { - pipeline.name = count == 5 ? "sd.rms_rope_cast.vulkan" : "sd.rms_rope_apply.vulkan"; - pipeline.spirv = count == 5 ? sd_vulkan_shaders::rms_rope_cast : sd_vulkan_shaders::rms_rope_apply; - pipeline.spirv_size = count == 5 ? sizeof(sd_vulkan_shaders::rms_rope_cast) : sizeof(sd_vulkan_shaders::rms_rope_apply); + 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; 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 index 7d08f74149..ef2f581c73 100644 --- a/src/core/vulkan/shaders/rms_rope_apply.comp +++ b/src/core/vulkan/shaders/rms_rope_apply.comp @@ -12,6 +12,7 @@ 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 { @@ -96,8 +97,10 @@ void main() { } } #ifdef CAST_F16 - value *= pc.output_scale; - value += 0.0; + 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; From 79526930040ba361c78cb42e79610cc5bfdc301e Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 6 Nov 2026 00:35:58 +0800 Subject: [PATCH 20/57] feat: support packed INT4 convrot models (#2276) Co-authored-by: noctrex --- docs/ggml_extensions.md | 10 + docs/int8_convrot.md | 41 +- docs/model_config.md | 2 + src/core/cpu/registration.cpp | 1 + src/core/cpu/registration.h | 1 + src/core/cpu/w4.cpp | 95 ++ src/core/cuda/CMakeLists.txt | 4 +- src/core/cuda/int8.cu | 70 +- src/core/cuda/registration.cu | 2 + src/core/cuda/w4.cu | 873 ++++++++++++++++++ src/core/ggml_extensions.cpp | 95 ++ src/core/ggml_extensions.h | 19 + src/core/vulkan/CMakeLists.txt | 3 + src/core/vulkan/int8.cpp | 74 +- .../vulkan/shaders/mul_mat_w4a4_convrot.comp | 189 ++++ .../vulkan/shaders/mul_mat_w4a8_convrot.comp | 244 +++++ .../shaders/quantize_i8_convrot_h64.comp | 115 +++ src/model/adapter/ip_adapter.hpp | 4 +- src/model/adapter/lora.hpp | 10 + src/model/common/block.hpp | 2 +- src/model/common/ggml_block.hpp | 55 +- src/model/diffusion/boogu.hpp | 4 +- src/model/diffusion/ernie_image.hpp | 2 +- src/model/diffusion/fibo2.hpp | 4 +- src/model/diffusion/flux.hpp | 2 +- src/model/diffusion/flux3.hpp | 6 +- src/model/diffusion/krea2.hpp | 10 +- src/model/diffusion/lens.hpp | 4 +- src/model/diffusion/lingbot_video.hpp | 6 +- src/model/diffusion/llada_image.hpp | 6 +- src/model/diffusion/ltxv.hpp | 8 +- src/model/diffusion/minimax_h3.hpp | 10 +- src/model/diffusion/minit2i.hpp | 6 +- src/model/diffusion/mmdit.hpp | 6 +- src/model/diffusion/moss_image.hpp | 8 +- src/model/diffusion/pixart.hpp | 4 +- src/model/diffusion/qwen_image_2_1.hpp | 4 +- src/model/diffusion/sefi_image.hpp | 2 +- src/model/diffusion/unet.hpp | 6 +- src/model/diffusion/z_image.hpp | 4 +- src/model/diffusion/z_image_l2p.hpp | 2 +- src/model/te/t5.hpp | 4 +- src/model/upscaler/h3_latent_upscaler.hpp | 2 +- src/model_io/safetensors_io.cpp | 58 +- src/model_io/tensor_storage.h | 5 + src/model_loader.cpp | 17 +- src/model_manager_files.cpp | 2 +- src/pipeline/diffusion_engine.cpp | 4 +- 48 files changed, 2003 insertions(+), 102 deletions(-) create mode 100644 src/core/cpu/w4.cpp create mode 100644 src/core/cuda/w4.cu create mode 100644 src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp create mode 100644 src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp create mode 100644 src/core/vulkan/shaders/quantize_i8_convrot_h64.comp diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 5c059892e8..77075989dd 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -104,6 +104,16 @@ and backend `graph_plan_*` interfaces do not dispatch these extensions. ## Backend services +Packed ComfyUI INT4 linear layers use `sd.linear_w4_convrot.v1`, built by +`ggml_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the +explicit output of `sd.quantize_i8_convrot.v1`; F32 callers get a quantization +node during graph construction. Logical output dimensions are stored in the +operator attributes, so backend copies of the packed activation do not need to +retain the producer's op marker. W4A4 uses row scales; W4A8 additionally exposes +the codebook, channel scales, and FP8 relative scales as tensor dependencies. +CPU, CUDA/HIP/MUSA, and Vulkan register kernels through the same extension API +as INT8. No ggml submodule changes are required. + Services are discovered through `ggml_backend_reg_get_proc_address`: | Service | Purpose | diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index e8036c60ea..49e77bcefc 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -100,7 +100,46 @@ The packed runtime activation tensor contains the I8 activation rows and their f 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 shared CUDA/HIP/MUSA convrot activation paths currently require a group size of `256`; other supported group sizes use CPU execution. +The shared CUDA/HIP/MUSA convrot activation paths support H64 and H256. Vulkan also supports H64 activation quantization; its INT8 matrix multiplication path currently requires H256. + +## Packed INT4 convrot + +ComfyUI `convrot_w4a4` and `asym_w4a8_int8` safetensors load directly, including +checkpoints that mix packed INT4, INT8 tensorwise, and floating-point layers. +Packed weights retain their checkpoint byte layout; they are not expanded or +converted to GGUF Q4 blocks at load time. + +- `convrot_w4a4`: signed 4-bit weights, low nibble first, with F32 per-output + `weight_scale`. Activations use the same INT8 convrot quantization as above; + the format name does not select a native INT4 activation kernel. +- `asym_w4a8_int8`: 4-bit indices into `weight_codebook` (16 F32 values), F32 + `weight_s_channel`, and F8_E4M3 `weight_s_rel` per 16 input features. The + metadata must specify `group_size: 16`. CPU evaluates the floating-point + codebook; GPU kernels requantize it to INT8 with `amax / 127`, introducing a + small additional rounding error. + +Both formats imply convrot rotation even without a `convrot` flag; missing +`convrot_groupsize` defaults to 256. + +Both formats reuse the convrot activation cache when the input scale is `1` and +produce F32 outputs with optional bias. CPU provides a portable implementation. +CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS; +CUDA and supported AMD targets use packed byte-dot instructions, with portable +byte arithmetic for other targets. Vulkan requires integer dot product support +and at least 13312 bytes of shared memory for W4A4 or 16452 bytes for W4A8. +GPU W4 kernels require the logical input width to be divisible by 64; H64 and +H256 are supported. Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA +by the active compilation target's +intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout. +There is no GPU architecture allowlist: newer targets supporting either +intrinsic automatically use its matching layout. Selection happens per device +image, including multi-architecture builds; targets supporting neither intrinsic +retain the integer-dot fallback. W4A4 retains integer dot products on HIP. + +LoRA application remains at runtime, using the original unrotated activations +and logical weight dimensions. Packed weights and their codebook/scale tensors +are excluded from weight type conversion. `SD_USE_UPSTREAM_GGML` builds reject +these formats because they require the application extension API. ## Example diff --git a/docs/model_config.md b/docs/model_config.md index 0d9290c1d5..d42c3226ce 100644 --- a/docs/model_config.md +++ b/docs/model_config.md @@ -60,6 +60,8 @@ static UNetConfig detect_from_weights(const String2TensorStorage& tensor_storage ``` Use `TensorStorage` metadata, especially `n_dims` and `ne`, to infer shapes. +Use `logical_ne0()` for Linear input widths; packed INT4 weights store half of +the logical width in `ne[0]`. Keep raw `ne` values for storage layout calculations. Do not load or parse tensor data for config detection. Detection should respect `prefix`. For nested weights, construct full names from diff --git a/src/core/cpu/registration.cpp b/src/core/cpu/registration.cpp index d63fd1f099..16e13dfe22 100644 --- a/src/core/cpu/registration.cpp +++ b/src/core/cpu/registration.cpp @@ -5,6 +5,7 @@ namespace sd_cpu { void register_extensions(ggml_backend_dev_t device) { register_i8(device); + register_w4(device); register_rope(device); register_modulate(device); } diff --git a/src/core/cpu/registration.h b/src/core/cpu/registration.h index eba629c2a6..77f4bacfb0 100644 --- a/src/core/cpu/registration.h +++ b/src/core/cpu/registration.h @@ -9,6 +9,7 @@ namespace sd_cpu { void register_rope(ggml_backend_dev_t device); void register_modulate(ggml_backend_dev_t device); void register_i8(ggml_backend_dev_t device); + void register_w4(ggml_backend_dev_t device); } #endif diff --git a/src/core/cpu/w4.cpp b/src/core/cpu/w4.cpp new file mode 100644 index 0000000000..4da0f2b51d --- /dev/null +++ b/src/core/cpu/w4.cpp @@ -0,0 +1,95 @@ +#include "core/cpu/registration.h" +#include "core/ggml_extensions.h" + +#ifndef SD_USE_UPSTREAM_GGML +#include + +namespace sd_cpu { + struct W4Execution { + SDW4Operation op; + float* output; + int64_t rows; + }; + + static void multiply_w4(ggml_tensor*, int ith, int nth, void* data) { + auto& e = *static_cast(data); + const auto& op = e.op; + const int64_t k = op.input->ne[0], n = op.weight->ne[1]; + const auto* w = static_cast(op.weight->data); + const auto* x = static_cast(op.input->data); + const auto* xs = reinterpret_cast(x + k * ((e.rows + 3) / 4 * 4)); + const auto* ws = static_cast(op.weight_scale->data); + const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; + const auto* sc = op.s_channel ? static_cast(op.s_channel->data) : nullptr; + const auto* sr = op.s_rel ? static_cast(op.s_rel->data) : nullptr; + const auto to_float = ggml_get_type_traits(GGML_TYPE_F8_E4M3)->to_float; + for (int64_t index = ith; index < n * e.rows; index += nth) { + const int64_t row = index / n, col = index % n; + const auto* packed = w + col * (k / 2); + float value; + if (op.kind == SDW4Kind::W4A4) { + uint32_t sum = 0; + for (int64_t i = 0; i < k; ++i) { + const int nibble = (packed[i / 2] >> (4 * (i % 2))) & 15; + sum += uint32_t(((nibble ^ 8) - 8) * int32_t(x[row * k + i])); + } + int32_t signed_sum; + std::memcpy(&signed_sum, &sum, sizeof(sum)); + value = float(signed_sum) * xs[row] * ws[col]; + } else { + float sum = 0.f; + for (int64_t g = 0; g < k / 16; ++g) { + float relative; + to_float(sr + col * (k / 16) + g, &relative, 1); + const float scale = sc[col] * relative; + for (int j = 0; j < 16; ++j) { + const int64_t i = g * 16 + j; + const int code = (packed[i / 2] >> (4 * (i % 2))) & 15; + sum += (ws[code] * scale) * float(x[row * k + i]); + } + } + value = sum * xs[row]; + } + e.output[index] = value + (bias ? bias[col] : 0.f); + } + } + + static bool w4_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) { + SDW4Operation op; + return sd_ggml_w4_describe(node, &op); + } + static ggml_status w4_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) { + *bytes = ggml_nbytes(nodes[0]); + return GGML_STATUS_SUCCESS; + } + static ggml_status w4_execute(const ggml_ext_execution_context* ctx, ggml_tensor* const* nodes, int, void*, void*) { + auto* dst = nodes[0]; + W4Execution e = {}; + if (!sd_ggml_w4_describe(dst, &e.op)) + return GGML_STATUS_FAILED; + auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(ctx->backend)); + auto run = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1")); + if (!run) + return GGML_STATUS_FAILED; + e.rows = ggml_nrows(dst); + e.output = static_cast(ggml_backend_buffer_get_base(ctx->workspace)); + auto status = run(ctx->backend, dst, multiply_w4, &e); + if (status == GGML_STATUS_SUCCESS) + std::memcpy(dst->data, e.output, ggml_nbytes(dst)); + return status; + } + void register_w4(ggml_backend_dev_t device) { + ggml_ext_kernel_desc d = {}; + d.abi_version = GGML_EXT_ABI_VERSION; + d.name = "sd.linear_w4_convrot.cpu.v1"; + d.op = GGML_OP_CUSTOM; + d.custom_op = sd_ggml_w4_op(); + d.priority = 100; + d.supports = w4_supports; + d.prepare = w4_prepare; + d.execute = w4_execute; + d.destroy = [](void*, void*) {}; + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + } +} +#endif diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 0a95d6cf7c..593445bf0e 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -2,7 +2,7 @@ cmake_policy(SET CMP0079 NEW) function(sd_cuda_sources VENDOR LANGUAGE) set(SOURCES) - foreach(OP rope int8 convrot affine modulate bf16 rms_affine registration) + foreach(OP rope int8 w4 convrot affine modulate bf16 rms_affine registration) set(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/${OP}_${VENDOR}.cu") string(TOUPPER "${VENDOR}" UPPER_VENDOR) file(GENERATE OUTPUT "${SOURCE}" CONTENT "#define SD_CUDA_VENDOR_${UPPER_VENDOR} 1\n#include \"core/cuda/${OP}.cu\"\n") @@ -49,6 +49,8 @@ if(SD_HIPBLAS) set_target_properties(${SD_CUDA_TARGET} PROPERTIES HIP_STANDARD 17 HIP_STANDARD_REQUIRED ON) endif() set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "-ffast-math") + # W4A8 WMMA uses the wave64 fragment layout on both GFX11 and GFX12. + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_hip.cu" APPEND PROPERTY COMPILE_OPTIONS "-mwavefrontsize64") target_link_libraries(${SD_CUDA_TARGET} PRIVATE hip::host roc::hipblas) target_link_libraries(${SD_LIB} PRIVATE hip::host roc::hipblas) endif() diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 45f6ecdf14..60f36bb1e9 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -386,12 +386,55 @@ namespace SD_CUDA_NAMESPACE { return info.int8_blas && ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1"); } + template + static __global__ void convrot_h64(const float* src, int8_t* dst, float* partial, const float* scales, int64_t k) { + const int64_t group = blockIdx.x, groups = k / 64, row = group / groups; + const int tid = threadIdx.x; + __shared__ float values[64]; + __shared__ float maxima[16]; +#pragma unroll + for (int j = 0; j < 4; ++j) + values[tid + j * 16] = src[group * 64 + tid + j * 16] * (1.f / 8.f); + __syncthreads(); +#pragma unroll + for (int stride = 1; stride < 64; stride *= 4) { + const int i = (tid / stride) * 4 * stride + tid % stride; + const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride]; + values[i] = convrot_h4_cuda(a, b, c, d, 0); + values[i + stride] = convrot_h4_cuda(a, b, c, d, 1); + values[i + 2 * stride] = convrot_h4_cuda(a, b, c, d, 2); + values[i + 3 * stride] = convrot_h4_cuda(a, b, c, d, 3); + __syncthreads(); + } + if constexpr (Quantize) { + const float inverse = scales[row] == 0.f ? 0.f : 1.f / scales[row]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int value = __float2int_rn(values[tid + j * 16] * inverse); + dst[group * 64 + tid + j * 16] = int8_t(max(-127, min(127, value))); + } + } else { + float amax = 0.f; +#pragma unroll + for (int j = 0; j < 4; ++j) + amax = fmaxf(amax, fabsf(values[tid + j * 16])); + maxima[tid] = amax; + __syncthreads(); + for (int stride = 8; stride; stride >>= 1) { + if (tid < stride) + maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]); + __syncthreads(); + } + if (!tid) + partial[group] = maxima[0]; + } + } static bool convrot_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { int group; - if (!sd_ggml_convrot_describe(node, &group) || group != 256 || !aligned(node) || !aligned(node->src[0])) + if (!sd_ggml_convrot_describe(node, &group) || (group != 1 && group != 64 && group != 256) || !aligned(node) || !aligned(node->src[0])) return false; const int64_t k = node->ne[0], rows = ggml_nrows(node->src[0]); - return k <= INT_MAX && rows <= INT_MAX && rows <= INT_MAX / (k / 256) && cuda_device_info(device).device >= 0; + return k <= INT_MAX && rows <= INT_MAX && (group == 1 || rows <= INT_MAX / (k / group)) && cuda_device_info(device).device >= 0; } enum Kind { MAT, QUANTIZE, @@ -421,7 +464,8 @@ namespace SD_CUDA_NAMESPACE { 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) && + int group; + if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) && size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) return 3; } @@ -430,6 +474,7 @@ namespace SD_CUDA_NAMESPACE { struct Plan { ggml_ext_gpu_device_info device; Kind kind; + int group; int64_t k, n, rows, padded; size_t quant, scales, accum, partial, output, gemm_offset; SDI8CudaGemmPlan* gemm = nullptr; @@ -456,7 +501,10 @@ namespace SD_CUDA_NAMESPACE { 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; + p->group = mat ? op.group_size : 0; + if (!mat) + sd_ggml_convrot_describe(dst, &p->group); + size_t total = 0; if (mat) { p->quant = reserve(total, op.input->type == GGML_TYPE_F32 || (!op.packed && p->rows != p->padded) ? size_t(p->k) * p->padded : 0); p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0); @@ -480,7 +528,8 @@ namespace SD_CUDA_NAMESPACE { 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->partial = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) + : 0); p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst)); *workspace = total; return GGML_STATUS_SUCCESS; @@ -506,6 +555,17 @@ namespace SD_CUDA_NAMESPACE { return error; } const size_t shared = size_t(p.k) * sizeof(float); + if (p.group == 1) { + quantize_rowwise_i8_cuda<<>>(src, dst, scales, p.k, p.rows); + return cuda_last_error(); + } + if (p.group == 64) { + const auto groups = p.k / 64; + convrot_h64<<>>(src, dst, partial, scales, p.k); + reduce_convrot_row_amax_i8_cuda<<>>(partial, scales, groups, p.rows); + convrot_h64<<>>(src, dst, partial, scales, p.k); + return cuda_last_error(); + } if (shared + 128 <= p.device.shared_memory) { static std::mutex mutex; static std::unordered_map configured; diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu index d75728021f..1f48174313 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -9,6 +9,7 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device); namespace SD_CUDA_NAMESPACE { void register_rope(ggml_backend_dev_t device); void register_i8(ggml_backend_dev_t device); + void register_w4(ggml_backend_dev_t device); void register_convrot(ggml_backend_dev_t device); void register_affine(ggml_backend_dev_t device); void register_modulate(ggml_backend_dev_t device); @@ -25,6 +26,7 @@ namespace SD_CUDA_NAMESPACE { register_affine(device); register_modulate(device); register_i8(device); + register_w4(device); register_convrot(device); register_rope(device); register_bf16(device); diff --git a/src/core/cuda/w4.cu b/src/core/cuda/w4.cu new file mode 100644 index 0000000000..ef46d02369 --- /dev/null +++ b/src/core/cuda/w4.cu @@ -0,0 +1,873 @@ +#include +#include "core/ggml_extensions.h" +#include "runtime.h" + +namespace SD_CUDA_NAMESPACE { + namespace sd_w4_cuda { + static __device__ __forceinline__ int shuffle_sum(int value, int lane, int width) { +#if defined(SD_CUDA_VENDOR_HIP) + // ROCm LLVM 24 can crash lowering __shfl's wave_shuffle on RDNA2. + // These 8/16-lane groups stay within one ds_bpermute half-wave. + const int index = (lane & (width - 1)) + (__lane_id() & ~(width - 1)); + return __builtin_amdgcn_ds_bpermute(index * 4, value); +#else + return cuda_shuffle(value, lane, width); +#endif + } + + // Signed byte dots preserve the INT4 grid after the XOR-8 bias correction. + static __device__ __forceinline__ int dot4(int a, int b, int c) { +#if defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__) || defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1200__) || defined(__gfx1201__)) + return __builtin_amdgcn_sudot4(true, a, true, b, c, false); +#elif defined(SD_CUDA_VENDOR_HIP) && (defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__) || defined(__gfx906__) || defined(__gfx908__) || defined(__gfx90a__) || defined(__gfx942__)) + return __builtin_amdgcn_sdot4(a, b, c, false); +#elif defined(SD_CUDA_VENDOR_MUSA) || (!defined(SD_CUDA_VENDOR_HIP) && __CUDA_ARCH__ >= 610) + return __dp4a(a, b, c); +#else + const auto* av = reinterpret_cast(&a); + const auto* bv = reinterpret_cast(&b); + for (int i = 0; i < 4; ++i) + c += int(av[i]) * int(bv[i]); + return c; +#endif + } + static __device__ __forceinline__ float fp8_e4m3(uint8_t v) { + const int exponent = (v >> 3) & 15, mantissa = v & 7; + float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9) + : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7); + return v & 128 ? -result : result; + } + // A CTA shares decoded weights across 32 activation rows. Each double-buffered + // stage covers 64 K values; W4A8 restores the codebook scale per 16-value group. + + static __device__ __forceinline__ uint32_t perm_bytes(uint32_t lo, uint32_t hi, uint32_t sel) { +#if defined(SD_CUDA_VENDOR_HIP) + return __builtin_amdgcn_perm(hi, lo, sel); +#elif !defined(SD_CUDA_VENDOR_MUSA) + return __byte_perm(lo, hi, sel); +#else + uint32_t r = 0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const uint32_t b = (sel >> (4 * i)) & 0xF; + const uint32_t byte = (b < 4 ? lo : hi) >> (8 * (b & 3)) & 0xFF; + r |= byte << (8 * i); + } + return r; +#endif + } + +#if defined(SD_CUDA_VENDOR_HIP) +#define PERM_SEL_EVEN 0x06040200u +#define PERM_SEL_ODD 0x07050301u +#else +#define PERM_SEL_EVEN 0x6420u +#define PERM_SEL_ODD 0x7531u +#endif + + static __device__ __forceinline__ void interleave_act_pair( + const int lo, + const int hi, + int* __restrict__ dst) { + dst[0] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_EVEN); + dst[1] = (int)perm_bytes((uint32_t)lo, (uint32_t)hi, PERM_SEL_ODD); + } + +#define MUL_MAT_W4_BLOCK 16 +#define MUL_MAT_W4_ROWS (2 * MUL_MAT_W4_BLOCK) +#define MUL_MAT_W4_GPS 4 +#define MUL_MAT_W4_SLOTS (4 * MUL_MAT_W4_GPS) +#define MUL_MAT_W4_OUT (MUL_MAT_W4_BLOCK * 4) + + static __global__ void __launch_bounds__(MUL_MAT_W4_BLOCK* MUL_MAT_W4_BLOCK) mul_mat_w4a4_convrot_dp4a( + const char* __restrict__ w, + const char* __restrict__ x, + int* __restrict__ y, + const uint32_t n, + const uint32_t rows, + const uint32_t packed_k) { + const uint32_t o_base = blockIdx.x * MUL_MAT_W4_OUT; + const uint32_t r0 = blockIdx.y * MUL_MAT_W4_ROWS + threadIdx.y; + const uint32_t r1 = r0 + MUL_MAT_W4_BLOCK; + const uint32_t lx = threadIdx.x; + const uint32_t ly = threadIdx.y; + const uint32_t idx = ly * MUL_MAT_W4_BLOCK + lx; + const uint32_t dc = idx & (MUL_MAT_W4_OUT - 1); + const uint32_t dw = idx >> 6; + + __shared__ int w_dec[2][MUL_MAT_W4_SLOTS * MUL_MAT_W4_OUT]; + __shared__ int x_tile[2][MUL_MAT_W4_ROWS * MUL_MAT_W4_SLOTS]; + + const uint32_t w_words = packed_k / 4; + const uint32_t x_words = (packed_k * 2) / 4; + const uint32_t groups = packed_k / 8; + const int* wi = (const int*)w; + const int* xi = (const int*)x; + + const uint32_t fl = lx & 7; + const uint32_t fr = (lx >> 3) ? r1 : r0; + const uint32_t tr = (lx >> 3) * MUL_MAT_W4_BLOCK + ly; + int my_sum = 0; + + { + const uint32_t ocol = o_base + dc; + int2 wv = make_int2(0, 0); + if (ocol < n && dw * 2 < w_words) { + wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + dw]; + } + w_dec[0][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808; + w_dec[0][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808; + w_dec[0][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808; + w_dec[0][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808; + + const uint32_t gi = fl >> 1; + const uint32_t xb = gi * 4 + (fl & 1) * 2; + const int a = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0; + const int b = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0; + interleave_act_pair(a, b, &x_tile[0][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]); + my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0); + } + __syncthreads(); + + int acc0a = 0, acc0b = 0; + int acc1a = 0, acc1b = 0; + int acc2a = 0, acc2b = 0; + int acc3a = 0, acc3b = 0; + int acc4a = 0, acc4b = 0; + int acc5a = 0, acc5b = 0; + int acc6a = 0, acc6b = 0; + int acc7a = 0, acc7b = 0; + + for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4_GPS) { + const uint32_t buf = (gs >> 2) & 1; + + if (gs + MUL_MAT_W4_GPS < groups) { + const uint32_t nbuf = buf ^ 1; + const uint32_t ngs = gs + MUL_MAT_W4_GPS; + { + const uint32_t ocol = o_base + dc; + const uint32_t widx = ngs * 2 + dw * 2; + int2 wv = make_int2(0, 0); + if (ocol < n && widx < w_words) { + wv = ((const int2*)wi)[(size_t)ocol * (w_words / 2) + ngs + dw]; + } + w_dec[nbuf][((dw * 2 + 0) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.x & 0x0F0F0F0F) ^ 0x08080808; + w_dec[nbuf][((dw * 2 + 0) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.x >> 4) & 0x0F0F0F0F) ^ 0x08080808; + w_dec[nbuf][((dw * 2 + 1) * 2 + 0) * MUL_MAT_W4_OUT + dc] = (wv.y & 0x0F0F0F0F) ^ 0x08080808; + w_dec[nbuf][((dw * 2 + 1) * 2 + 1) * MUL_MAT_W4_OUT + dc] = ((wv.y >> 4) & 0x0F0F0F0F) ^ 0x08080808; + const uint32_t gi = ngs + (fl >> 1); + const uint32_t xb = gi * 4 + (fl & 1) * 2; + const int a = (fr < rows && xb < x_words) ? (int)xi[(size_t)fr * x_words + xb] : 0; + const int b = (fr < rows && xb + 1 < x_words) ? (int)xi[(size_t)fr * x_words + xb + 1] : 0; + interleave_act_pair(a, b, &x_tile[nbuf][tr * MUL_MAT_W4_SLOTS + (fl >> 1) * 4 + (fl & 1) * 2]); + my_sum += dot4(a, 0x01010101, 0) + dot4(b, 0x01010101, 0); + } + } +#pragma unroll + for (uint32_t gl = 0; gl < MUL_MAT_W4_GPS; ++gl) { + const int u0[4] = { + x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 0], + x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 1], + x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 2], + x_tile[buf][ly * MUL_MAT_W4_SLOTS + gl * 4 + 3], + }; + const int u1[4] = { + x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 0], + x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 1], + x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 2], + x_tile[buf][(MUL_MAT_W4_BLOCK + ly) * MUL_MAT_W4_SLOTS + gl * 4 + 3], + }; +#pragma unroll + for (uint32_t jp = 0; jp < 2; ++jp) { + const int ue0 = jp == 0 ? u0[0] : u0[2]; + const int uo0 = jp == 0 ? u0[1] : u0[3]; + const int ue1 = jp == 0 ? u1[0] : u1[2]; + const int uo1 = jp == 0 ? u1[1] : u1[3]; + const uint32_t we = (gl * 4 + jp * 2) * MUL_MAT_W4_OUT; + const uint32_t wo = we + MUL_MAT_W4_OUT; + + acc0a = dot4(w_dec[buf][we + lx], ue0, acc0a); + acc1a = dot4(w_dec[buf][we + 16 + lx], ue0, acc1a); + acc2a = dot4(w_dec[buf][we + 32 + lx], ue0, acc2a); + acc3a = dot4(w_dec[buf][we + 48 + lx], ue0, acc3a); + acc0b = dot4(w_dec[buf][wo + lx], uo0, acc0b); + acc1b = dot4(w_dec[buf][wo + 16 + lx], uo0, acc1b); + acc2b = dot4(w_dec[buf][wo + 32 + lx], uo0, acc2b); + acc3b = dot4(w_dec[buf][wo + 48 + lx], uo0, acc3b); + acc4a = dot4(w_dec[buf][we + lx], ue1, acc4a); + acc5a = dot4(w_dec[buf][we + 16 + lx], ue1, acc5a); + acc6a = dot4(w_dec[buf][we + 32 + lx], ue1, acc6a); + acc7a = dot4(w_dec[buf][we + 48 + lx], ue1, acc7a); + acc4b = dot4(w_dec[buf][wo + lx], uo1, acc4b); + acc5b = dot4(w_dec[buf][wo + 16 + lx], uo1, acc5b); + acc6b = dot4(w_dec[buf][wo + 32 + lx], uo1, acc6b); + acc7b = dot4(w_dec[buf][wo + 48 + lx], uo1, acc7b); + } + } + __syncthreads(); + } + + int v = my_sum; + v += shuffle_sum(v, (int(threadIdx.x) % 8 + 4) % 8, 8); + v += shuffle_sum(v, (int(threadIdx.x) % 8 + 2) % 8, 8); + v += shuffle_sum(v, (int(threadIdx.x) % 8 + 1) % 8, 8); + const int sum0 = shuffle_sum(v, 0, 16); + const int sum1 = shuffle_sum(v, 8, 16); + + const int accv[2][4] = { + {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b}, + {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b}, + }; + const int sumv[2] = {sum0, sum1}; +#pragma unroll + for (uint32_t t = 0; t < 2; ++t) { + const uint32_t rt = t == 0 ? r0 : r1; + if (rt < rows) { +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const uint32_t oj = o_base + j * 16 + lx; + if (oj < n) { + y[(size_t)rt * n + oj] = accv[t][j] - 8 * sumv[t]; + } + } + } + } + } + +#define MUL_MAT_W4A8_BLOCK 16 +#define MUL_MAT_W4A8_ROWS (2 * MUL_MAT_W4A8_BLOCK) +#define MUL_MAT_W4A8_GPS 4 +#define MUL_MAT_W4A8_SLOTS (4 * MUL_MAT_W4A8_GPS) +#define MUL_MAT_W4A8_OUT (MUL_MAT_W4A8_BLOCK * 4) + +#if defined(SD_CUDA_VENDOR_HIP) && defined(__HIP_DEVICE_COMPILE__) +// __has_builtin checks the active device target's features, not just whether +// the compiler recognizes the name. Prefer the non-replicated fragment layout. +#if __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12) +#define SD_W4_WMMA_GFX12 +#elif __has_builtin(__builtin_amdgcn_wmma_i32_16x16x16_iu8_w64) +#define SD_W4_WMMA_GFX11 +#endif +#endif + +// Keep each intrinsic paired with its fragment layout, including on newer GPUs. +#if defined(SD_W4_WMMA_GFX11) + typedef int v4si32_w4a8 __attribute__((ext_vector_type(4))); + + static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma( + const char* __restrict__ w, + const char* __restrict__ x, + const float* __restrict__ codebook, + const float* __restrict__ s_channel, + const uint8_t* __restrict__ s_rel, + const float* __restrict__ act_scales, + const float* __restrict__ bias, + float* __restrict__ y, + const uint32_t n, + const uint32_t rows, + const uint32_t k, + const uint32_t r_base) { + const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT; + const uint32_t idx = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x; + const uint32_t wv = idx >> 6; + const uint32_t lane = idx & 63; + const uint32_t j = lane & 15; + const uint32_t q = lane >> 4; + const uint32_t ocol = o_base + wv * 16 + j; + const uint32_t ra = r_base + j; + + __shared__ __align__(4) int8_t cb8[16]; + __shared__ float s_cb; + __shared__ int lut[256]; + __shared__ float s_lut[256]; + + const uint32_t w_words = k / 8; + const uint32_t x_words = k / 4; + const uint32_t groups = k / 16; + const int* wi = (const int*)w; + const int* xi = (const int*)x; + + if (idx == 0) { + float amax = 0.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + amax = fmaxf(amax, fabsf(codebook[i])); + } + s_cb = amax / 127.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb); + cb8[i] = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v)); + } +#pragma unroll + for (int b = 0; b < 256; ++b) { + lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8); + } +#pragma unroll + for (int b = 0; b < 256; ++b) { + s_lut[b] = fp8_e4m3(uint8_t(b)); + } + } + __syncthreads(); + + const bool o_valid = ocol < n; + const float s_chan = o_valid ? s_channel[ocol] : 0.0f; + const float scb = s_cb; + + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + const v4si32_w4a8 zero = {}; + + int w0c = 0, w1c = 0, xc0 = 0, xc1 = 0, xc2 = 0, xc3 = 0, sc = 0; + int w0n = 0, w1n = 0, xn0 = 0, xn1 = 0, xn2 = 0, xn3 = 0, sn = 0; + { + const int* xr = xi + (size_t)ra * x_words; + xc0 = ra < rows ? xr[0] : 0; + xc1 = ra < rows ? xr[1] : 0; + xc2 = ra < rows ? xr[2] : 0; + xc3 = ra < rows ? xr[3] : 0; + if (o_valid) { + const int* wp = wi + (size_t)ocol * w_words; + w0c = wp[0]; + w1c = wp[1]; + sc = s_rel[(size_t)ocol * groups]; + } + } + if (groups > 1) { + const int* xr = xi + (size_t)ra * x_words + 4; + xn0 = ra < rows ? xr[0] : 0; + xn1 = ra < rows ? xr[1] : 0; + xn2 = ra < rows ? xr[2] : 0; + xn3 = ra < rows ? xr[3] : 0; + if (o_valid) { + const int* wp = wi + (size_t)ocol * w_words + 2; + w0n = wp[0]; + w1n = wp[1]; + sn = s_rel[(size_t)ocol * groups + 1]; + } + } + + for (uint32_t g = 0; g < groups; ++g) { + v4si32_w4a8 bfrag = {}; + float s = 0.0f; + if (o_valid) { + bfrag = { + (lut[w0c & 0xFF] & 0xFFFF) | ((lut[(w0c >> 8) & 0xFF] & 0xFFFF) << 16), + (lut[(w0c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w0c >> 24) & 0xFF] & 0xFFFF) << 16), + (lut[w1c & 0xFF] & 0xFFFF) | ((lut[(w1c >> 8) & 0xFF] & 0xFFFF) << 16), + (lut[(w1c >> 16) & 0xFF] & 0xFFFF) | ((lut[(w1c >> 24) & 0xFF] & 0xFFFF) << 16), + }; + + s = scb * (s_chan * s_lut[sc & 0xFF]); + } + const v4si32_w4a8 afrag = {xc0, xc1, xc2, xc3}; + const v4si32_w4a8 dfrag = + __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64(true, afrag, true, bfrag, zero, false); + + w0c = w0n; + w1c = w1n; + sc = sn; + xc0 = xn0; + xc1 = xn1; + xc2 = xn2; + xc3 = xn3; + if (g + 2 < groups) { + const int* xr = xi + (size_t)ra * x_words + (g + 2) * 4; + xn0 = ra < rows ? xr[0] : 0; + xn1 = ra < rows ? xr[1] : 0; + xn2 = ra < rows ? xr[2] : 0; + xn3 = ra < rows ? xr[3] : 0; + if (o_valid) { + const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2; + w0n = wp[0]; + w1n = wp[1]; + sn = s_rel[(size_t)ocol * groups + g + 2]; + } + } + + acc[0] = fmaf((float)dfrag.x, s, acc[0]); + acc[1] = fmaf((float)dfrag.y, s, acc[1]); + acc[2] = fmaf((float)dfrag.z, s, acc[2]); + acc[3] = fmaf((float)dfrag.w, s, acc[3]); + } + + if (o_valid) { +#pragma unroll + for (uint32_t d = 0; d < 4; ++d) { + const uint32_t r = r_base + 4 * d + q; + if (r < rows) { + float value = acc[d] * act_scales[r]; + if (bias != nullptr) { + value += bias[ocol]; + } + y[(size_t)r * n + ocol] = value; + } + } + } + } +#endif + +#if defined(SD_W4_WMMA_GFX12) + typedef int v4si32_w4a8_g12 __attribute__((ext_vector_type(4))); + + static __device__ __forceinline__ void mul_mat_w4a8_convrot_wmma_gfx12( + const char* __restrict__ w, + const char* __restrict__ x, + const float* __restrict__ codebook, + const float* __restrict__ s_channel, + const uint8_t* __restrict__ s_rel, + const float* __restrict__ act_scales, + const float* __restrict__ bias, + float* __restrict__ y, + const uint32_t n, + const uint32_t rows, + const uint32_t k, + const uint32_t r_base) { + const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT; + const uint32_t idx = threadIdx.y * MUL_MAT_W4A8_BLOCK + threadIdx.x; + const uint32_t wv = idx >> 6; + const uint32_t lane = idx & 63; + const uint32_t j = lane & 15; + const uint32_t h = lane >> 4; + const uint32_t kh = ((h & 1) << 1) | (h >> 1); + const uint32_t ocol = o_base + wv * 16 + j; + const uint32_t ra = r_base + j; + + __shared__ __align__(4) int8_t cb8[16]; + __shared__ float s_cb; + __shared__ int lut[256]; + __shared__ float s_lut[256]; + + const uint32_t w_words = k / 8; + const uint32_t x_words = k / 4; + const uint32_t groups = k / 16; + const int* wi = (const int*)w; + const int* xi = (const int*)x; + + if (idx == 0) { + float amax = 0.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + amax = fmaxf(amax, fabsf(codebook[i])); + } + s_cb = amax / 127.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int v = s_cb == 0.f ? 0 : lrintf(codebook[i] / s_cb); + cb8[i] = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v)); + } +#pragma unroll + for (int b = 0; b < 256; ++b) { + lut[b] = (cb8[b & 0xF] & 0xFF) | ((cb8[b >> 4] & 0xFF) << 8); + } +#pragma unroll + for (int b = 0; b < 256; ++b) { + s_lut[b] = fp8_e4m3(uint8_t(b)); + } + } + __syncthreads(); + + const bool o_valid = ocol < n; + const float s_chan = o_valid ? s_channel[ocol] : 0.0f; + const float scb = s_cb; + + float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + const v4si32_w4a8_g12 zero = {}; + + int w0c = 0, w1c = 0, xc = 0, sc = 0; + int w0n = 0, w1n = 0, xn = 0, sn = 0; + { + const int* xr = xi + (size_t)ra * x_words; + xc = ra < rows ? xr[kh] : 0; + if (o_valid) { + const int* wp = wi + (size_t)ocol * w_words; + w0c = wp[0]; + w1c = wp[1]; + sc = s_rel[(size_t)ocol * groups]; + } + } + if (groups > 1) { + const int* xr = xi + (size_t)ra * x_words; + xn = ra < rows ? xr[4 + kh] : 0; + if (o_valid) { + const int* wp = wi + (size_t)ocol * w_words + 2; + w0n = wp[0]; + w1n = wp[1]; + sn = s_rel[(size_t)ocol * groups + 1]; + } + } + + for (uint32_t g = 0; g < groups; ++g) { + int bfrag = 0; + float s = 0.0f; + if (o_valid) { + const int bwd = (kh & 2) ? w1c : w0c; + const int bsh = (kh & 1) * 16; + bfrag = (lut[(bwd >> bsh) & 0xFF] & 0xFFFF) | + ((lut[(bwd >> (bsh + 8)) & 0xFF] & 0xFFFF) << 16); + + s = scb * (s_chan * s_lut[sc & 0xFF]); + } + const v4si32_w4a8_g12 dfrag = + __builtin_amdgcn_wmma_i32_16x16x16_iu8_w64_gfx12(true, xc, true, bfrag, zero, false); + + w0c = w0n; + w1c = w1n; + sc = sn; + xc = xn; + if (g + 2 < groups) { + const int* xr = xi + (size_t)ra * x_words; + xn = ra < rows ? xr[(g + 2) * 4 + kh] : 0; + if (o_valid) { + const int* wp = wi + (size_t)ocol * w_words + (g + 2) * 2; + w0n = wp[0]; + w1n = wp[1]; + sn = s_rel[(size_t)ocol * groups + g + 2]; + } + } + + acc[0] = fmaf((float)dfrag.x, s, acc[0]); + acc[1] = fmaf((float)dfrag.y, s, acc[1]); + acc[2] = fmaf((float)dfrag.z, s, acc[2]); + acc[3] = fmaf((float)dfrag.w, s, acc[3]); + } + + if (o_valid) { +#pragma unroll + for (uint32_t d = 0; d < 4; ++d) { + const uint32_t r = r_base + 4 * kh + d; + if (r < rows) { + float value = acc[d] * act_scales[r]; + if (bias != nullptr) { + value += bias[ocol]; + } + y[(size_t)r * n + ocol] = value; + } + } + } + } +#endif + + static __global__ void __launch_bounds__(MUL_MAT_W4A8_BLOCK* MUL_MAT_W4A8_BLOCK) mul_mat_w4a8_convrot( + const char* __restrict__ w, + const char* __restrict__ x, + const float* __restrict__ codebook, + const float* __restrict__ s_channel, + const uint8_t* __restrict__ s_rel, + const float* __restrict__ act_scales, + const float* __restrict__ bias, + float* __restrict__ y, + const uint32_t n, + const uint32_t rows, + const uint32_t k) { +#if defined(SD_W4_WMMA_GFX11) || defined(SD_W4_WMMA_GFX12) + // Retain the 32-row launch grid shared with DP4A. Each WMMA wave + // computes a 16-row tile; the barrier protects the reused LUTs. + for (uint32_t tile = 0; tile < 2; ++tile) { +#if defined(SD_W4_WMMA_GFX12) + mul_mat_w4a8_convrot_wmma_gfx12(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k, + blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK); +#else + mul_mat_w4a8_convrot_wmma(w, x, codebook, s_channel, s_rel, act_scales, bias, y, n, rows, k, + blockIdx.y * MUL_MAT_W4A8_ROWS + tile * MUL_MAT_W4A8_BLOCK); +#endif + __syncthreads(); + } +#else + const uint32_t o_base = blockIdx.x * MUL_MAT_W4A8_OUT; + const uint32_t r0 = blockIdx.y * MUL_MAT_W4A8_ROWS + threadIdx.y; + const uint32_t r1 = r0 + MUL_MAT_W4A8_BLOCK; + const uint32_t lx = threadIdx.x; + const uint32_t ly = threadIdx.y; + const uint32_t idx = ly * MUL_MAT_W4A8_BLOCK + lx; + const uint32_t dc = idx & (MUL_MAT_W4A8_OUT - 1); + const uint32_t dw = idx >> 6; + + __shared__ int w_dec[2][MUL_MAT_W4A8_SLOTS * MUL_MAT_W4A8_OUT]; + __shared__ int x_tile[2][MUL_MAT_W4A8_ROWS * MUL_MAT_W4A8_SLOTS]; + __shared__ float s_dec[2][MUL_MAT_W4A8_GPS * MUL_MAT_W4A8_OUT]; + __shared__ __align__(4) int8_t cb8[16]; + __shared__ float s_cb; + __shared__ int lut[256]; + + const uint32_t w_words = k / 8; + const uint32_t x_words = k / 4; + const uint32_t groups = k / 16; + const int* wi = (const int*)w; + const int* xi = (const int*)x; + + if (ly == 0 && lx == 0) { + float amax = 0.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + amax = fmaxf(amax, fabsf(codebook[i])); + } + + s_cb = amax / 127.0f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int v = s_cb == 0.f ? 0 : __float2int_rn(codebook[i] / s_cb); + cb8[i] = (int8_t)(v < -127 ? -127 : (v > 127 ? 127 : v)); + } + } + __syncthreads(); + + lut[idx] = (cb8[idx & 0xF] & 0xFF) | ((cb8[idx >> 4] & 0xFF) << 8); + __syncthreads(); + +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + const uint32_t pw = dw * 2 + j; + const uint32_t ocol = o_base + dc; + const int wv = (ocol < n && pw < w_words) ? wi[(size_t)ocol * w_words + pw] : 0; + w_dec[0][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16); + w_dec[0][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16); + } + { + const uint32_t ocol = o_base + dc; + float s = 0.0f; + if (ocol < n && dw < groups) { + s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + dw]); + } + s_dec[0][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s; + } + if (lx < 16) { + const uint32_t aidx = lx; + x_tile[0][ly * MUL_MAT_W4A8_SLOTS + lx] = + (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0; + x_tile[0][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] = + (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0; + } + __syncthreads(); + + float acc0 = 0.0f; + float acc1 = 0.0f; + float acc2 = 0.0f; + float acc3 = 0.0f; + float acc4 = 0.0f; + float acc5 = 0.0f; + float acc6 = 0.0f; + float acc7 = 0.0f; + + for (uint32_t gs = 0; gs < groups; gs += MUL_MAT_W4A8_GPS) { + const uint32_t buf = (gs >> 2) & 1; + + if (gs + MUL_MAT_W4A8_GPS < groups) { + const uint32_t nbuf = buf ^ 1; + const uint32_t ngs = gs + MUL_MAT_W4A8_GPS; +#pragma unroll + for (uint32_t j = 0; j < 2; ++j) { + const uint32_t pw = dw * 2 + j; + const uint32_t ocol = o_base + dc; + const uint32_t widx = ngs * 2 + pw; + const int wv = (ocol < n && widx < w_words) ? wi[(size_t)ocol * w_words + widx] : 0; + w_dec[nbuf][(pw * 2 + 0) * MUL_MAT_W4A8_OUT + dc] = (lut[wv & 0xFF] & 0xFFFF) | ((lut[(wv >> 8) & 0xFF] & 0xFFFF) << 16); + w_dec[nbuf][(pw * 2 + 1) * MUL_MAT_W4A8_OUT + dc] = (lut[(wv >> 16) & 0xFF] & 0xFFFF) | ((lut[(wv >> 24) & 0xFF] & 0xFFFF) << 16); + } + { + const uint32_t ocol = o_base + dc; + float s = 0.0f; + if (ocol < n && ngs + dw < groups) { + s = s_channel[ocol] * fp8_e4m3(s_rel[(size_t)ocol * groups + ngs + dw]); + } + s_dec[nbuf][dw * MUL_MAT_W4A8_OUT + dc] = s_cb * s; + } + if (lx < 16) { + const uint32_t aidx = ngs * 4 + lx; + x_tile[nbuf][ly * MUL_MAT_W4A8_SLOTS + lx] = + (r0 < rows && aidx < x_words) ? xi[(size_t)r0 * x_words + aidx] : 0; + x_tile[nbuf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + lx] = + (r1 < rows && aidx < x_words) ? xi[(size_t)r1 * x_words + aidx] : 0; + } + } + +#pragma unroll + for (uint32_t gl = 0; gl < MUL_MAT_W4A8_GPS; ++gl) { + const int u0[4] = { + x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 0], + x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 1], + x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 2], + x_tile[buf][ly * MUL_MAT_W4A8_SLOTS + gl * 4 + 3], + }; + const int u1[4] = { + x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 0], + x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 1], + x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 2], + x_tile[buf][(MUL_MAT_W4A8_BLOCK + ly) * MUL_MAT_W4A8_SLOTS + gl * 4 + 3], + }; + int accg0 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u0[0], 0); + accg0 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u0[1], accg0); + accg0 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u0[2], accg0); + accg0 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u0[3], accg0); + int accg1 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u0[0], 0); + accg1 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u0[1], accg1); + accg1 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u0[2], accg1); + accg1 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u0[3], accg1); + int accg2 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u0[0], 0); + accg2 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u0[1], accg2); + accg2 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u0[2], accg2); + accg2 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u0[3], accg2); + int accg3 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u0[0], 0); + accg3 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u0[1], accg3); + accg3 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u0[2], accg3); + accg3 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u0[3], accg3); + int accg4 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + lx], u1[0], 0); + accg4 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + lx], u1[1], accg4); + accg4 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + lx], u1[2], accg4); + accg4 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + lx], u1[3], accg4); + int accg5 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 16 + lx], u1[0], 0); + accg5 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 16 + lx], u1[1], accg5); + accg5 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 16 + lx], u1[2], accg5); + accg5 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 16 + lx], u1[3], accg5); + int accg6 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 32 + lx], u1[0], 0); + accg6 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 32 + lx], u1[1], accg6); + accg6 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 32 + lx], u1[2], accg6); + accg6 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 32 + lx], u1[3], accg6); + int accg7 = dot4(w_dec[buf][(gl * 4 + 0) * MUL_MAT_W4A8_OUT + 48 + lx], u1[0], 0); + accg7 = dot4(w_dec[buf][(gl * 4 + 1) * MUL_MAT_W4A8_OUT + 48 + lx], u1[1], accg7); + accg7 = dot4(w_dec[buf][(gl * 4 + 2) * MUL_MAT_W4A8_OUT + 48 + lx], u1[2], accg7); + accg7 = dot4(w_dec[buf][(gl * 4 + 3) * MUL_MAT_W4A8_OUT + 48 + lx], u1[3], accg7); + acc0 = fmaf((float)accg0, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc0); + acc1 = fmaf((float)accg1, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc1); + acc2 = fmaf((float)accg2, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc2); + acc3 = fmaf((float)accg3, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc3); + acc4 = fmaf((float)accg4, s_dec[buf][gl * MUL_MAT_W4A8_OUT + lx], acc4); + acc5 = fmaf((float)accg5, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 16 + lx], acc5); + acc6 = fmaf((float)accg6, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 32 + lx], acc6); + acc7 = fmaf((float)accg7, s_dec[buf][gl * MUL_MAT_W4A8_OUT + 48 + lx], acc7); + } + __syncthreads(); + } + + const float accv[8] = {acc0, acc1, acc2, acc3, acc4, acc5, acc6, acc7}; +#pragma unroll + for (uint32_t t = 0; t < 2; ++t) { + const uint32_t rt = t == 0 ? r0 : r1; + if (rt < rows) { +#pragma unroll + for (uint32_t j = 0; j < 4; ++j) { + const uint32_t oj = o_base + j * 16 + lx; + if (oj < n) { + float value = accv[t * 4 + j] * act_scales[rt]; + if (bias != nullptr) { + value += bias[oj]; + } + y[(size_t)rt * n + oj] = value; + } + } + } + } +#endif + } + + static __global__ void dequantize(const int32_t* accum, const float* xs, const float* ws, const float* bias, float* dst, int64_t n, int64_t total) { + const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < total) { + float value = float(accum[i]) * xs[i / n] * ws[i % n]; + dst[i] = value + (bias ? bias[i % n] : 0.f); + } + } + static bool aligned(const ggml_tensor* t, size_t alignment) { + if (!t) + return true; + size_t offset = 0; + while (!t->data && t->view_src) { + offset += t->view_offs; + t = t->view_src; + } + return (reinterpret_cast(t->data) + offset) % alignment == 0; + } + static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + SDW4Operation op; + if (!sd_ggml_w4_describe(node, &op) || cuda_device_info(device).device < 0) + return false; + const int64_t k = op.input->ne[0], n = node->ne[0], rows = ggml_nrows(node); + if (k % 64 || k > INT_MAX || n > INT_MAX - 63 || rows > 65535 * 32 || + ggml_nelements(node) > int64_t(INT_MAX) * 256 || !aligned(op.weight, 8) || !aligned(node, 4)) + return false; + for (auto* t : node->src) + if (!aligned(t, 4)) + return false; + return true; + } + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) { + SDW4Operation op; + if (!sd_ggml_w4_describe(nodes[0], &op)) + return GGML_STATUS_FAILED; + const auto* dst = nodes[0]; + bool may_alias = dst->data || dst->view_src; + if (dst->data && dst->buffer && !dst->view_src) { + may_alias = false; + for (const auto* input : dst->src) { + if (!input) + continue; + size_t offset = 0; + const size_t size = ggml_nbytes(input); + while (!input->data && input->view_src) { + offset += input->view_offs; + input = input->view_src; + } + if (!input->data) + continue; + const auto a = reinterpret_cast(dst->data), b = reinterpret_cast(input->data) + offset; + may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst)); + } + } + *bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 ? 1 : 0) + (may_alias ? 1 : 0)); + return GGML_STATUS_SUCCESS; + } + 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 ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) { + auto* dst = nodes[0]; + SDW4Operation op; + ggml_ext_gpu_context context; + if (!sd_ggml_w4_describe(dst, &op) || !cuda_context(e->backend, &context)) + return GGML_STATUS_FAILED; + const auto stream = static_cast(context.stream); + const int64_t n = dst->ne[0], rows = ggml_nrows(dst), k = op.input->ne[0]; + auto* scratch = e->workspace ? static_cast(ggml_backend_buffer_get_base(e->workspace)) : nullptr; + bool alias = false; + for (auto* input : dst->src) + alias |= overlaps(dst, input); + auto* output = alias ? scratch : static_cast(dst->data); + const auto* w = static_cast(op.weight->data); + const auto* x = static_cast(op.input->data); + const auto* xs = reinterpret_cast(x + k * ((rows + 3) / 4 * 4)); + const auto* ws = static_cast(op.weight_scale->data); + const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; + const dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16); + if (op.kind == SDW4Kind::W4A4) { + auto* accum = reinterpret_cast(scratch + (alias ? ggml_nelements(dst) : 0)); + mul_mat_w4a4_convrot_dp4a<<>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2)); + if (cuda_last_error() != cuda_success) + return GGML_STATUS_FAILED; + dequantize<<>>(accum, xs, ws, bias, output, n, n * rows); + } else { + mul_mat_w4a8_convrot<<>>(w, x, ws, + static_cast(op.s_channel->data), static_cast(op.s_rel->data), + xs, bias, output, uint32_t(n), uint32_t(rows), uint32_t(k)); + } + auto error = cuda_last_error(); + if (error == cuda_success && alias) + error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream); + return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; + } + } + void register_w4(ggml_backend_dev_t device) { + ggml_ext_kernel_desc d = {}; + d.abi_version = GGML_EXT_ABI_VERSION; + d.name = "sd.linear_w4_convrot." SD_CUDA_NAME ".v1"; + d.op = GGML_OP_CUSTOM; + d.custom_op = sd_ggml_w4_op(); + d.priority = 100; + d.supports = sd_w4_cuda::supports; + d.prepare = sd_w4_cuda::prepare; + d.execute = sd_w4_cuda::execute; + d.destroy = [](void*, void*) {}; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); + } +} diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index 27089f6c3b..25fa22265f 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -301,6 +301,101 @@ ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int gro return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr); } +struct W4Attributes { + int64_t shape[3]; + SDW4Kind kind; + int32_t has_bias; +}; + +static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { + if (!attributes || count < 3) + return false; + W4Attributes attr; + std::memcpy(&attr, attributes, sizeof(attr)); + const bool codebook = attr.kind == SDW4Kind::W4A8; + if ((!codebook && attr.kind != SDW4Kind::W4A4) || (attr.has_bias != 0 && attr.has_bias != 1) || + count != (codebook ? 5 : 3) + attr.has_bias) + return false; + for (int i = 0; i < count; ++i) + if (!inputs[i] || !ggml_is_contiguous(inputs[i])) + return false; + const auto* w = inputs[0]; + const auto* x = inputs[1]; + if (w->type != GGML_TYPE_I8 || w->ne[0] <= 0 || w->ne[0] > i8_size_limit / 2 || + w->ne[1] <= 0 || w->ne[2] != 1 || w->ne[3] != 1) + return false; + const int64_t k = 2 * w->ne[0], n = w->ne[1]; + int64_t rows = 1; + for (auto dim : attr.shape) { + if (dim <= 0 || dim > i8_size_limit / rows) + return false; + rows *= dim; + } + int64_t height; + if (!packed_height(k, rows, &height) || n > i8_size_limit / 4 / rows || + x->type != GGML_TYPE_I8 || x->ne[0] != k || x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1 || + !scale_matches(inputs[2], codebook ? 16 : n) || + (attr.has_bias && !scale_matches(inputs[count - 1], n))) + return false; + if (codebook && (k % 16 || !scale_matches(inputs[3], n) || inputs[4]->type != GGML_TYPE_F8_E4M3 || + inputs[4]->ne[0] != k / 16 || inputs[4]->ne[1] != n || inputs[4]->ne[2] != 1 || inputs[4]->ne[3] != 1)) + return false; + out->type = GGML_TYPE_F32; + out->ne[0] = n; + std::memcpy(out->ne + 1, attr.shape, sizeof(attr.shape)); + return true; +} + +ggml_ext_op_t sd_ggml_w4_op() { + static const auto op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(W4Attributes), w4_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) { + if (!node || ggml_ext_op_get(node) != sd_ggml_w4_op() || !ggml_is_contiguous(node)) + return false; + W4Attributes attr; + std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr)); + const bool codebook = attr.kind == SDW4Kind::W4A8; + const int count = (codebook ? 5 : 3) + (attr.has_bias != 0); + ggml_ext_tensor_desc out; + if (!w4_infer(node->src, count, &attr, &out) || node->type != out.type || + std::memcmp(node->ne, out.ne, sizeof(out.ne))) + return false; + *op = {node->src[0], node->src[1], node->src[2], codebook ? node->src[3] : nullptr, + codebook ? node->src[4] : nullptr, attr.has_bias ? node->src[count - 1] : nullptr, attr.kind}; + return true; +} + +ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* sc, ggml_tensor* sr, ggml_tensor* bias, SDW4Kind kind, int group) { + if (!w || !x || group < 0 || (group && !valid_group(group))) + return nullptr; + const int quant_group = group ? group : 1; + if (x->type == GGML_TYPE_F32) + x = ggml_quantize_i8_convrot(ctx, x, quant_group); + int packed_group; + if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group) + return nullptr; + const auto* logical = x->src[0]; + W4Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0}; + ggml_tensor* inputs[6] = {w, x, ws}; + int count = 3; + if (kind == SDW4Kind::W4A8) { + inputs[count++] = sc; + inputs[count++] = sr; + } else if (sc || sr) { + return nullptr; + } + if (bias) + inputs[count++] = bias; + return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &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; diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index 86c9121a7a..bc9ee9a5c4 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -39,6 +39,25 @@ 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 class SDW4Kind : int32_t { W4A4 = 1, + W4A8 = 2 }; + +struct SDW4Operation { + const ggml_tensor* weight; + const ggml_tensor* input; + const ggml_tensor* weight_scale; + const ggml_tensor* s_channel; + const ggml_tensor* s_rel; + const ggml_tensor* bias; + SDW4Kind kind; +}; + +ggml_ext_op_t sd_ggml_w4_op(); +bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation); +// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook +// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node. +ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* s_channel, ggml_tensor* s_rel, ggml_tensor* bias, SDW4Kind kind, int group_size); + enum ggml_sage_attn_mode { GGML_SAGE_ATTN_AUTO = 0, GGML_SAGE_ATTN_2, diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt index 55acf0bb91..e71bcd5339 100644 --- a/src/core/vulkan/CMakeLists.txt +++ b/src/core/vulkan/CMakeLists.txt @@ -53,6 +53,9 @@ 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_h64 quantize_i8_convrot_h64 ${SD_VK_DOT} -O -DGROUP_SIZE=64) +sd_vulkan_shader(mul_mat_w4a4_convrot mul_mat_w4a4_convrot ${SD_VK_DOT} -O) +sd_vulkan_shader(mul_mat_w4a8_convrot mul_mat_w4a8_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) diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index 7e3e67d48d..5ac8151854 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -11,7 +11,10 @@ #include "mul_mat_i8_tensorwise.h" #include "mul_mat_i8_tensorwise_cm1.h" #include "mul_mat_i8_tensorwise_cm2.h" +#include "mul_mat_w4a4_convrot.h" +#include "mul_mat_w4a8_convrot.h" #include "quantize_i8_convrot.h" +#include "quantize_i8_convrot_h64.h" #include "quantize_i8_convrot_register.h" #include "quantize_i8_convrot_subgroup.h" #include "rms_modulate_convrot.h" @@ -111,6 +114,23 @@ namespace sd_vulkan_i8 { (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 w4_supports(ggml_backend_dev_t dev, const ggml_tensor* dst, void*) { + Device d; + SDW4Operation op; + if (!device_info(dev, d) || !sd_ggml_w4_describe(dst, &op) || !bounded(d, dst) || + sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot) <= 4 || sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot) <= 4 || + d.compute.shared_memory < (op.kind == SDW4Kind::W4A8 ? 16452u : 13312u) || d.info.max_push_constants < 16) + return false; + for (auto* t : dst->src) + if (t && (!bounded(d, t) || ggml_nbytes(t) > UINT32_MAX)) + return false; + const uint64_t k = op.input->ne[0], n = dst->ne[0], rows = ggml_nrows(dst); + return k % 64 == 0 && k <= UINT32_MAX && n <= UINT32_MAX && rows <= UINT32_MAX && + ggml_nelements(dst) <= UINT32_MAX && + (k * ((rows + 3) / 4 * 4)) % d.info.buffer_alignment == 0 && + (n + 63) / 64 <= d.info.max_group_count[0] && (rows + 31) / 32 <= 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; @@ -534,9 +554,12 @@ namespace sd_vulkan_i8 { 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; + int group; + if (!sd_ggml_convrot_describe(output, &group)) + return GGML_STATUS_FAILED; + auto pipeline = group == 64 ? SD_VK_PIPELINE(quantize_i8_convrot_h64, 3, QuantParameters) : SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters); + bool subgroup = group == 256 && 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) { @@ -653,6 +676,42 @@ namespace sd_vulkan_i8 { } return finish(e, d, output, dst, alias, status); } + static ggml_status multiply_w4(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* output) { + SDW4Operation op; + if (!sd_ggml_w4_describe(output, &op)) + return GGML_STATUS_FAILED; + const uint32_t k = uint32_t(op.input->ne[0]), n = uint32_t(output->ne[0]), rows = uint32_t(ggml_nrows(output)); + const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4); + const ggml_ext_vulkan_binding xs = {op.input, nullptr, scale_offset, size_t(rows) * 4}; + std::vector bindings = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}}; + if (op.kind == SDW4Kind::W4A8) { + bindings.push_back({op.s_channel}); + bindings.push_back({op.s_rel}); + } + bindings.push_back(xs); + bindings.push_back({op.bias ? op.bias : op.weight_scale}); + ggml_ext_vulkan_binding dst; + bool alias; + if (!output_binding(e, d, output, bindings, dst, alias)) + return GGML_STATUS_FAILED; + bindings.push_back(dst); + struct Parameters { + uint32_t n, rows, k, bias; + }; + const Parameters pc = {n, rows, k, op.bias ? 1u : 0u}; + auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters); + const uint32_t groups[] = {(n + 63) / 64, (rows + 31) / 32, 1}; + auto status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups); + return finish(e, d, output, dst, alias, status); + } + static bool quant64_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) { + Device d; + int group; + return device_info(dev, d) && sd_ggml_convrot_describe(n, &group) && group == 64 && + bounded(d, n) && bounded(d, n->src[0]) && n->ne[0] <= UINT32_MAX && + ggml_nrows(n->src[0]) <= UINT32_MAX && d.info.max_group_invocations >= 16 && d.info.max_group_size[0] >= 16 && + (uint64_t(n->ne[0]) * ((ggml_nrows(n->src[0]) + 3) / 4 * 4)) % d.info.buffer_alignment == 0; + } 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); @@ -747,6 +806,8 @@ namespace sd_vulkan_i8 { Device d; if (!device_info(ggml_backend_get_device(e->backend), d)) return GGML_STATUS_FAILED; + if (ggml_ext_op_get(nodes[0]) == sd_ggml_w4_op()) + return multiply_w4(e, d, nodes[0]); return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count); } #undef SD_VK_PIPELINE @@ -768,6 +829,13 @@ void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) { d.custom_op = sd_ggml_convrot_op(); d.supports = quant_supports; GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + d.name = "sd.quantize_i8_convrot.h64.vulkan.v1"; + d.supports = quant64_supports; + GGML_ASSERT(ggml_ext_register_kernel(device, &d)); + d.name = "sd.linear_w4_convrot.vulkan.v1"; + d.custom_op = sd_ggml_w4_op(); + d.supports = w4_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(); diff --git a/src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp b/src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp new file mode 100644 index 0000000000..f5fb10e6f4 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a4_convrot.comp @@ -0,0 +1,189 @@ +#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; +} +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 w_dec[2][16 * 64]; +shared int x_tile[2][32 * 16]; +shared int sum_p[2][16 * 8]; + +void interleave_pair(int a, int b, uint buf, uint li) { + const uint a0 = uint(a); + const uint b0 = uint(b); + x_tile[buf][li + 0] = int((a0 & 0x000000FFu) | ((a0 >> 8) & 0x0000FF00u) | ((b0 << 16) & 0x00FF0000u) | ((b0 << 8) & 0xFF000000u)); + x_tile[buf][li + 1] = int(((a0 >> 8) & 0x000000FFu) | ((a0 >> 16) & 0x0000FF00u) | ((b0 << 8) & 0x00FF0000u) | (b0 & 0xFF000000u)); +} + +void main() { + const uint out_base = gl_WorkGroupID.x * 64u; + const uint lx = gl_LocalInvocationID.x; + const uint ly = gl_LocalInvocationID.y; + const uint idx = gl_LocalInvocationIndex; + const uint w_words = p.k / 8u; + const uint a_words = p.k / 4u; + const uint groups = p.k / 16u; + + const uint dc = idx & 63u; + const uint dw = idx >> 6u; + const uint ocol = out_base + dc; + + const uint fl = lx & 7u; + const uint ft = lx >> 3u; + const uint fr = ft == 0u ? gl_WorkGroupID.y * 32u + ly : gl_WorkGroupID.y * 32u + ly + 16u; + const uint tr = ft * 16u + ly; + int my_sum = 0; + int acc0a = 0; + int acc0b = 0; + int acc1a = 0; + int acc1b = 0; + int acc2a = 0; + int acc2b = 0; + int acc3a = 0; + int acc3b = 0; + int acc4a = 0; + int acc4b = 0; + int acc5a = 0; + int acc5b = 0; + int acc6a = 0; + int acc6b = 0; + int acc7a = 0; + int acc7b = 0; + + { + [[unroll]] for (uint j = 0; j < 2; ++j) { + const uint pw = dw * 2 + j; + const uint widx = pw; + const uint wv = ocol < p.n && widx < w_words + ? uint(data_w[ocol * w_words + widx]) + : 0u; + w_dec[0][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u); + w_dec[0][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u); + } + const uint gi = fl >> 1u; + const uint xb = gi * 4u + (fl & 1u) * 2u; + const int a = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0; + const int b = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0; + interleave_pair(a, b, 0, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u); + my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101))); + } + barrier(); + + for (uint gs = 0; gs < groups; gs += 4) { + const uint buf = (gs >> 2) & 1u; + + if (gs + 4 < groups) { + const uint nbuf = buf ^ 1u; + const uint ngs = gs + 4; + [[unroll]] for (uint j = 0; j < 2; ++j) { + const uint pw = dw * 2 + j; + const uint widx = ngs * 2 + pw; + const uint wv = ocol < p.n && widx < w_words + ? uint(data_w[ocol * w_words + widx]) + : 0u; + w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((wv & 0x0F0F0F0Fu) ^ 0x08080808u); + w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int(((wv >> 4) & 0x0F0F0F0Fu) ^ 0x08080808u); + } + const uint gi = ngs + (fl >> 1u); + const uint xb = gi * 4u + (fl & 1u) * 2u; + const int a = fr < p.rows && xb < a_words ? data_a[fr * a_words + xb] : 0; + const int b = fr < p.rows && xb + 1u < a_words ? data_a[fr * a_words + xb + 1u] : 0; + interleave_pair(a, b, nbuf, tr * 16u + (fl >> 1u) * 4u + (fl & 1u) * 2u); + my_sum += int(dotPacked4x8EXT(a, int(0x01010101))) + int(dotPacked4x8EXT(b, int(0x01010101))); + } + + [[unroll]] for (uint gl = 0; gl < 4; ++gl) { + const int u0e = x_tile[buf][ly * 16u + gl * 4u + 0u]; + const int u0o = x_tile[buf][ly * 16u + gl * 4u + 1u]; + const int u0e2 = x_tile[buf][ly * 16u + gl * 4u + 2u]; + const int u0o2 = x_tile[buf][ly * 16u + gl * 4u + 3u]; + const int u1e = x_tile[buf][(16u + ly) * 16u + gl * 4u + 0u]; + const int u1o = x_tile[buf][(16u + ly) * 16u + gl * 4u + 1u]; + const int u1e2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 2u]; + const int u1o2 = x_tile[buf][(16u + ly) * 16u + gl * 4u + 3u]; + [[unroll]] for (uint jp = 0; jp < 2; ++jp) { + const int ue0 = jp == 0u ? u0e : u0e2; + const int uo0 = jp == 0u ? u0o : u0o2; + const int ue1 = jp == 0u ? u1e : u1e2; + const int uo1 = jp == 0u ? u1o : u1o2; + const uint we = (gl * 4u + jp * 2u) * 64u; + const uint wo = we + 64u; + acc0a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue0)); + acc1a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue0)); + acc2a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue0)); + acc3a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue0)); + acc0b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo0)); + acc1b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo0)); + acc2b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo0)); + acc3b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo0)); + acc4a += int(dotPacked4x8EXT(w_dec[buf][we + lx], ue1)); + acc5a += int(dotPacked4x8EXT(w_dec[buf][we + 16u + lx], ue1)); + acc6a += int(dotPacked4x8EXT(w_dec[buf][we + 32u + lx], ue1)); + acc7a += int(dotPacked4x8EXT(w_dec[buf][we + 48u + lx], ue1)); + acc4b += int(dotPacked4x8EXT(w_dec[buf][wo + lx], uo1)); + acc5b += int(dotPacked4x8EXT(w_dec[buf][wo + 16u + lx], uo1)); + acc6b += int(dotPacked4x8EXT(w_dec[buf][wo + 32u + lx], uo1)); + acc7b += int(dotPacked4x8EXT(w_dec[buf][wo + 48u + lx], uo1)); + } + } + barrier(); + } + + sum_p[ft][ly * 8u + fl] = my_sum; + barrier(); + int sum0 = 0; + int sum1 = 0; + [[unroll]] for (uint i = 0; i < 8; ++i) { + sum0 += sum_p[0][ly * 8u + i]; + sum1 += sum_p[1][ly * 8u + i]; + } + + const int accv[2][4] = { + {acc0a + acc0b, acc1a + acc1b, acc2a + acc2b, acc3a + acc3b}, + {acc4a + acc4b, acc5a + acc5b, acc6a + acc6b, acc7a + acc7b}, + }; + const int sumv[2] = {sum0, sum1}; + [[unroll]] for (uint t = 0; t < 2u; ++t) { + const uint rt = gl_WorkGroupID.y * 32u + ly + t * 16u; + if (rt < p.rows) { + [[unroll]] for (uint j = 0; j < 4; ++j) { + const uint oj = out_base + j * 16u + lx; + if (oj < p.n) { + float value = float(accv[t][j] - 8 * sumv[t]) * data_as[rt] * data_ws[oj]; + if (p.has_bias != 0) { + value += data_b[oj]; + } + data_d[rt * p.n + oj] = value; + } + } + } + } +} diff --git a/src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp b/src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp new file mode 100644 index 0000000000..4438ce2876 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a8_convrot.comp @@ -0,0 +1,244 @@ +#version 460 + +#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; +} +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 CB { + float data_cb[]; +}; +layout(binding = 3, std430) readonly buffer SC { + float data_sc[]; +}; +layout(binding = 4, std430) readonly buffer SR { + uint data_sr[]; +}; +layout(binding = 5, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 6, std430) readonly buffer B { + float data_b[]; +}; +layout(binding = 7, std430) writeonly buffer D { + float data_d[]; +}; + +shared int w_dec[2][16 * 64]; +shared int x_tile[2][32 * 16]; +shared float s_dec[2][4 * 64]; +shared int lut[256]; +shared float s_lut[256]; +shared int cb8s[16]; +shared float s_cb; + +float f8_e4m3_to_f32(uint v) { + const uint sign = (v >> 7) & 1u; + const uint exp = (v >> 3) & 0xFu; + const uint man = v & 0x7u; + float result; + if (exp == 0xFu && man == 0x7u) { + result = uintBitsToFloat(0x7FC00000u); + } else if (exp == 0u) { + result = ldexp(float(man), -9); + } else { + result = ldexp(1.0 + float(man) / 8.0, int(exp) - 7); + } + return sign != 0u ? -result : result; +} + +void main() { + const uint out_base = gl_WorkGroupID.x * 64u; + const uint lx = gl_LocalInvocationID.x; + const uint ly = gl_LocalInvocationID.y; + + const uint row0 = gl_WorkGroupID.y * 32u + ly; + const uint row1 = row0 + 16u; + const uint idx = gl_LocalInvocationIndex; + const uint w_words = p.k / 8u; + const uint a_words = p.k / 4u; + const uint groups = p.k / 16u; + const uint sr_words = groups / 4u; + + const uint dc = idx & 63u; + const uint dw = idx >> 6u; + const uint ocol = out_base + dc; + float sum0 = 0.0; + float sum1 = 0.0; + float sum2 = 0.0; + float sum3 = 0.0; + float sum4 = 0.0; + float sum5 = 0.0; + float sum6 = 0.0; + float sum7 = 0.0; + + if (idx == 0u) { + float amax = 0.0; + for (uint i = 0; i < 16; ++i) { + amax = max(amax, abs(data_cb[i])); + } + s_cb = amax / 127.0; + for (uint i = 0; i < 16; ++i) { + const int v = (s_cb == 0.0 ? 0 : int(roundEven(data_cb[i] / s_cb))); + cb8s[i] = v < -127 ? -127 : (v > 127 ? 127 : v); + } + } + barrier(); + + lut[idx] = (cb8s[idx & 0xFu] & 0xFF) | (cb8s[idx >> 4] << 8); + s_lut[idx] = f8_e4m3_to_f32(idx); + barrier(); + + { + [[unroll]] for (uint j = 0; j < 2; ++j) { + const uint pw = dw * 2 + j; + const uint widx = pw; + const uint wv = ocol < p.n && widx < w_words + ? uint(data_w[ocol * w_words + widx]) + : 0u; + + w_dec[0][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16)); + w_dec[0][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16)); + } + { + const uint g = dw; + float s = 0.0; + if (ocol < p.n && g < groups) { + const uint sr_word = data_sr[ocol * sr_words + g / 4u]; + s = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu]; + } + s_dec[0][dw * 64 + dc] = s_cb * s; + } + if (lx < 16u) { + const uint aidx = lx; + x_tile[0][ly * 16 + lx] = row0 < p.rows && aidx < a_words + ? data_a[row0 * a_words + aidx] + : 0; + x_tile[0][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words + ? data_a[row1 * a_words + aidx] + : 0; + } + } + barrier(); + + for (uint gs = 0; gs < groups; gs += 4) { + const uint buf = (gs >> 2) & 1u; + + if (gs + 4 < groups) { + const uint nbuf = buf ^ 1u; + const uint ngs = gs + 4; + [[unroll]] for (uint j = 0; j < 2; ++j) { + const uint pw = dw * 2 + j; + const uint widx = ngs * 2 + pw; + const uint wv = ocol < p.n && widx < w_words + ? uint(data_w[ocol * w_words + widx]) + : 0u; + w_dec[nbuf][(pw * 2 + 0) * 64 + dc] = int((uint(lut[wv & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 8) & 0xFFu] & 0xFFFF) << 16)); + w_dec[nbuf][(pw * 2 + 1) * 64 + dc] = int((uint(lut[(wv >> 16) & 0xFFu] & 0xFFFF)) | (uint(lut[(wv >> 24) & 0xFFu] & 0xFFFF) << 16)); + } + { + const uint g = ngs + dw; + float s = 0.0; + if (ocol < p.n && g < groups) { + const uint sr_word = data_sr[ocol * sr_words + g / 4u]; + s = data_sc[ocol] * s_lut[(sr_word >> (8u * (g % 4u))) & 0xFFu]; + } + s_dec[nbuf][dw * 64 + dc] = s_cb * s; + } + if (lx < 16u) { + const uint aidx = ngs * 4 + lx; + x_tile[nbuf][ly * 16 + lx] = row0 < p.rows && aidx < a_words + ? data_a[row0 * a_words + aidx] + : 0; + x_tile[nbuf][(16u + ly) * 16 + lx] = row1 < p.rows && aidx < a_words + ? data_a[row1 * a_words + aidx] + : 0; + } + } + + float partial0 = 0.0; + float partial1 = 0.0; + float partial2 = 0.0; + float partial3 = 0.0; + float partial4 = 0.0; + float partial5 = 0.0; + float partial6 = 0.0; + float partial7 = 0.0; + { + [[unroll]] for (uint gl = 0; gl < 4; ++gl) { + int accg0 = 0; + int accg1 = 0; + int accg2 = 0; + int accg3 = 0; + int accg4 = 0; + int accg5 = 0; + int accg6 = 0; + int accg7 = 0; + [[unroll]] for (uint q = 0; q < 4; ++q) { + const int a0 = x_tile[buf][ly * 16 + gl * 4 + q]; + const int a1 = x_tile[buf][(16u + ly) * 16 + gl * 4 + q]; + accg0 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a0)); + accg1 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a0)); + accg2 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a0)); + accg3 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a0)); + accg4 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + lx], a1)); + accg5 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 16u + lx], a1)); + accg6 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 32u + lx], a1)); + accg7 += int(dotPacked4x8EXT(w_dec[buf][(gl * 4 + q) * 64 + 48u + lx], a1)); + } + const float s0 = s_dec[buf][gl * 64 + lx]; + const float s1 = s_dec[buf][gl * 64 + 16u + lx]; + const float s2 = s_dec[buf][gl * 64 + 32u + lx]; + const float s3 = s_dec[buf][gl * 64 + 48u + lx]; + partial0 += float(accg0) * s0; + partial1 += float(accg1) * s1; + partial2 += float(accg2) * s2; + partial3 += float(accg3) * s3; + partial4 += float(accg4) * s0; + partial5 += float(accg5) * s1; + partial6 += float(accg6) * s2; + partial7 += float(accg7) * s3; + } + } + sum0 += partial0; + sum1 += partial1; + sum2 += partial2; + sum3 += partial3; + sum4 += partial4; + sum5 += partial5; + sum6 += partial6; + sum7 += partial7; + barrier(); + } + + const float sumv[8] = {sum0, sum1, sum2, sum3, sum4, sum5, sum6, sum7}; + [[unroll]] for (uint t = 0; t < 2u; ++t) { + const uint rt = t == 0u ? row0 : row1; + if (rt < p.rows) { + [[unroll]] for (uint j = 0; j < 4; ++j) { + const uint oj = out_base + j * 16u + lx; + if (oj < p.n) { + float value = sumv[t * 4u + j] * data_as[rt]; + if (p.has_bias != 0) { + value += data_b[oj]; + } + data_d[rt * p.n + oj] = value; + } + } + } + } +} diff --git a/src/core/vulkan/shaders/quantize_i8_convrot_h64.comp b/src/core/vulkan/shaders/quantize_i8_convrot_h64.comp new file mode 100644 index 0000000000..2fd6dbb463 --- /dev/null +++ b/src/core/vulkan/shaders/quantize_i8_convrot_h64.comp @@ -0,0 +1,115 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require + +#ifndef GROUP_SIZE +#define GROUP_SIZE 256 +#endif +#if GROUP_SIZE != 64 && GROUP_SIZE != 256 +#error unsupported GROUP_SIZE +#endif +#define GROUP_THREADS (GROUP_SIZE / 4) + +#if GROUP_SIZE == 256 +#define HADAMARD_NORM (1.0 / 16.0) +#else +#define HADAMARD_NORM (1.0 / 8.0) +#endif + +layout(local_size_x = GROUP_THREADS, 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[GROUP_SIZE]; +shared float maxima[GROUP_THREADS]; + +void transform_group(const uint row, const uint group) { + const uint tid = gl_LocalInvocationID.x; + const uint base_offset = row * p.k + group * GROUP_SIZE; + + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint index = tid + i * GROUP_THREADS; + values[index] = data_a[base_offset + index] * HADAMARD_NORM; + } + barrier(); + + [[unroll]] for (uint stride = 1; stride < GROUP_SIZE; 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 / GROUP_SIZE; + 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 * GROUP_THREADS])); + } + barrier(); + } + + maxima[tid] = local_max; + barrier(); + [[unroll]] for (uint stride = GROUP_THREADS / 2; 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 * GROUP_SIZE) / 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/model/adapter/ip_adapter.hpp b/src/model/adapter/ip_adapter.hpp index 9f8dc9f64f..b36977df4e 100644 --- a/src/model/adapter/ip_adapter.hpp +++ b/src/model/adapter/ip_adapter.hpp @@ -141,7 +141,7 @@ namespace IPAdapter { } auto proj_in_iter = tensor_storage_map.find(prefix + ".image_proj.proj_in.weight"); if (proj_in_iter != tensor_storage_map.end()) { - embed_dim = proj_in_iter->second.ne[0]; + embed_dim = proj_in_iter->second.logical_ne0(); } auto proj_out_iter = tensor_storage_map.find(prefix + ".image_proj.proj_out.weight"); if (proj_out_iter != tensor_storage_map.end()) { @@ -168,7 +168,7 @@ namespace IPAdapter { } auto proj_iter = tensor_storage_map.find(prefix + ".image_proj.proj.weight"); if (proj_iter != tensor_storage_map.end()) { - clip_dim = proj_iter->second.ne[0]; + clip_dim = proj_iter->second.logical_ne0(); out_dim = proj_iter->second.ne[1]; } num_tokens = out_dim / ctx_dim; diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index b2657a3691..7b9e8ab82d 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -1127,6 +1127,16 @@ struct MultiLoraAdapter : public WeightAdapter { ggml_tensor* output, const std::string& prefix, WeightAdapter::ForwardParams forward_params) override { + // Packed INT4 stores half the logical input width. Adapter shape checks + // and reshapes need the logical weight dimensions, without decoding it. + ggml_tensor logical_weight = *w; + if (forward_params.op_type == WeightAdapter::ForwardParams::op_type_t::OP_LINEAR && + w->type == GGML_TYPE_I8 && x->ne[0] == 2 * w->ne[0]) { + logical_weight.ne[0] = x->ne[0]; + for (int i = 1; i < GGML_MAX_DIMS; ++i) + logical_weight.nb[i] = logical_weight.nb[i - 1] * logical_weight.ne[i - 1]; + w = &logical_weight; + } for (auto& lora_model : lora_models) { ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false); if (weight_diff != nullptr) { diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp index e79ee2f12a..66180fa8d1 100644 --- a/src/model/common/block.hpp +++ b/src/model/common/block.hpp @@ -319,7 +319,7 @@ class CrossAttention : public GGMLBlock { tensor_storage_map.find(prefix + "to_k_ip.weight") != tensor_storage_map.end()) { has_ip = true; int64_t inner_dim = d_head * n_head; - int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").ne[0]; + int64_t ip_dim = tensor_storage_map.at(prefix + "to_k_ip.weight").logical_ne0(); blocks["to_k_ip"] = std::shared_ptr(new Linear(ip_dim, inner_dim, false)); blocks["to_v_ip"] = std::shared_ptr(new Linear(ip_dim, inner_dim, false)); } diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 7529848477..0653ae1fff 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -148,6 +148,7 @@ class Linear : public UnaryBlock { bool force_f32; bool force_prec_f32; bool has_weight_scale = false; + int w4_convrot_kind = 0; bool int8_convrot = false; int int8_convrot_group_size = 0; float scale; @@ -158,16 +159,21 @@ class Linear : public UnaryBlock { has_weight_scale = false; int8_convrot = false; int8_convrot_group_size = 0; + auto weight_storage = tensor_storage_map.find(prefix + "weight"); + w4_convrot_kind = weight_storage != tensor_storage_map.end() ? weight_storage->second.w4_convrot_kind : 0; enum ggml_type wtype = get_type(prefix + "weight", tensor_storage_map, GGML_TYPE_F32); - if (in_features % ggml_blck_size(wtype) != 0 || force_f32) { + if (!w4_convrot_kind && (in_features % ggml_blck_size(wtype) != 0 || force_f32)) { wtype = GGML_TYPE_F32; } - params["weight"] = ggml_new_tensor_2d(ctx, wtype, in_features, out_features); + if (w4_convrot_kind) { + GGML_ASSERT(wtype == GGML_TYPE_I8 && in_features % 2 == 0); + GGML_ASSERT(weight_storage->second.ne[0] == in_features / 2 && weight_storage->second.ne[1] == out_features); + } + params["weight"] = ggml_new_tensor_2d(ctx, wtype, w4_convrot_kind ? in_features / 2 : in_features, out_features); if (bias) { enum ggml_type wtype = GGML_TYPE_F32; params["bias"] = ggml_new_tensor_1d(ctx, wtype, out_features); } - auto weight_storage = tensor_storage_map.find(prefix + "weight"); const bool is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise; auto weight_scale_storage = tensor_storage_map.find(prefix + "weight_scale"); if (weight_scale_storage != tensor_storage_map.end()) { @@ -176,9 +182,16 @@ class Linear : public UnaryBlock { params["weight_scale"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, scale_nelements); has_weight_scale = true; } - if (is_int8_tensorwise) { + if (w4_convrot_kind == 2) { + params["weight_codebook"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 16); + params["weight_s_channel"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, out_features); +#ifndef SD_USE_UPSTREAM_GGML + params["weight_s_rel"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F8_E4M3, in_features / 16, out_features); +#endif + } + if (is_int8_tensorwise || w4_convrot_kind) { GGML_ASSERT(wtype == GGML_TYPE_I8); - GGML_ASSERT(has_weight_scale); + GGML_ASSERT(w4_convrot_kind == 2 || has_weight_scale); int8_convrot = weight_storage->second.int8_convrot; int8_convrot_group_size = weight_storage->second.int8_convrot_group_size; } @@ -259,13 +272,31 @@ class Linear : public UnaryBlock { } } #endif - out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx, - x, - w, - weight_scale, - b, - int8_convrot ? int8_convrot_group_size : 0, - scale); +#ifndef SD_USE_UPSTREAM_GGML + if (w4_convrot_kind) { + if (scale != 1.f) + x = ggml_ext_scale(ctx->ggml_ctx, x, scale); + out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x, + w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale, + w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr, + w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr, + scale == 1.f ? b : nullptr, static_cast(w4_convrot_kind), + int8_convrot ? int8_convrot_group_size : 0); + GGML_ASSERT(out); + if (scale != 1.f) { + out = ggml_ext_scale(ctx->ggml_ctx, out, 1.f / scale); + if (b) + out = ggml_add_inplace(ctx->ggml_ctx, out, b); + } + } else +#endif + out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx, + x, + w, + weight_scale, + b, + int8_convrot ? int8_convrot_group_size : 0, + scale); if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index 3f652e1589..ecdb7def3e 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -72,10 +72,10 @@ namespace Boogu { } if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { int64_t patch_area = config.patch_size * config.patch_size; - config.in_channels = tensor_storage.ne[0] / patch_area; + config.in_channels = tensor_storage.logical_ne0() / patch_area; config.hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "time_caption_embed.caption_embedder.1.weight") && tensor_storage.n_dims == 2) { - config.instruction_feat_dim = tensor_storage.ne[0]; + config.instruction_feat_dim = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "single_stream_layers.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) { detected_head_dim = tensor_storage.ne[0]; diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 15badb9834..5898ac0679 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -42,7 +42,7 @@ namespace ErnieImage { config.in_channels = tensor_storage.ne[2]; config.hidden_size = tensor_storage.ne[3]; } else if (ends_with(name, "text_proj.weight") && tensor_storage.n_dims == 2) { - config.text_in_dim = tensor_storage.ne[0]; + config.text_in_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "layers.0.self_attention.norm_q.weight")) { detected_head_dim = tensor_storage.ne[0]; } else if (ends_with(name, "layers.0.mlp.gate_proj.weight") && tensor_storage.n_dims == 2) { diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp index 0c42235f8f..2a65ea1d9c 100644 --- a/src/model/diffusion/fibo2.hpp +++ b/src/model/diffusion/fibo2.hpp @@ -31,12 +31,12 @@ namespace Fibo2 { } auto key = name.substr(prefix.size() + 1); if (key == "x_embedder.weight") { - config.channels = tensor.ne[0]; + config.channels = tensor.logical_ne0(); 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.text_dim = tensor.logical_ne0(); 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]; diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 40edca0ed9..843e70b458 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -158,7 +158,7 @@ namespace Flux { } } if (ends_with(name, "txt_in.weight")) { - config.context_in_dim = tensor_storage.ne[0]; + config.context_in_dim = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1]; } if (ends_with(name, "single_blocks.0.norm.key_norm.scale")) { diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp index fac1b1a567..105b22565d 100644 --- a/src/model/diffusion/flux3.hpp +++ b/src/model/diffusion/flux3.hpp @@ -27,12 +27,12 @@ namespace Flux3 { std::string key = name.substr(prefix.size() + 1); if (key == "txt_in.weight") { detected = true; - config.context_dim = tensor.ne[0]; + config.context_dim = tensor.logical_ne0(); config.hidden_size = tensor.ne[1]; } else if (key == "emb_in.image.weight") { - config.channels = tensor.ne[0]; + config.channels = tensor.logical_ne0(); } else if (key == "single_blocks.0.mlp_out.weight") { - config.mlp_dim = tensor.ne[0]; + config.mlp_dim = tensor.logical_ne0(); } else if (key == "ref_time_in.in_layer.weight") { config.flow_map = true; } diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index fba817aa67..e44c5a9c6e 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -92,11 +92,11 @@ namespace Krea2 { continue; } if (ends_with(name, "first.weight") && tensor_storage.n_dims == 2) { - config.in_channels = tensor_storage.ne[0] / (config.patch_size * config.patch_size); + config.in_channels = tensor_storage.logical_ne0() / (config.patch_size * config.patch_size); config.out_channels = config.in_channels; config.features = tensor_storage.ne[1]; } else if (ends_with(name, "blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) { - detected_head_dim = tensor_storage.ne[0]; + detected_head_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) { if (detected_head_dim > 0) { config.heads = tensor_storage.ne[1] / detected_head_dim; @@ -106,11 +106,11 @@ namespace Krea2 { config.kv_heads = tensor_storage.ne[1] / detected_head_dim; } } else if (ends_with(name, "txtfusion.projector.weight") && tensor_storage.n_dims == 2) { - config.text_layers = tensor_storage.ne[0]; + config.text_layers = tensor_storage.logical_ne0(); } else if (ends_with(name, "txtfusion.layerwise_blocks.0.prenorm.scale") && tensor_storage.n_dims == 1) { - config.text_dim = tensor_storage.ne[0]; + config.text_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.qknorm.qnorm.scale") && tensor_storage.n_dims == 1) { - detected_text_head_dim = tensor_storage.ne[0]; + detected_text_head_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "txtfusion.layerwise_blocks.0.attn.wq.weight") && tensor_storage.n_dims == 2) { if (detected_text_head_dim > 0) { config.text_heads = tensor_storage.ne[1] / detected_text_head_dim; diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index bcfefd97bf..0f8c816739 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -36,13 +36,13 @@ namespace Lens { continue; } if (ends_with(name, "img_in.weight") && tensor_storage.n_dims == 2) { - config.in_channels = tensor_storage.ne[0]; + config.in_channels = tensor_storage.logical_ne0(); int64_t inner_dim = tensor_storage.ne[1]; if (config.attention_head_dim > 0) { config.num_attention_heads = inner_dim / config.attention_head_dim; } } else if (ends_with(name, "txt_in.weight") && tensor_storage.n_dims == 2) { - config.selected_layer_count = static_cast(tensor_storage.ne[0] / config.joint_attention_dim); + config.selected_layer_count = static_cast(tensor_storage.logical_ne0() / config.joint_attention_dim); } else if (ends_with(name, "proj_out.weight") && tensor_storage.n_dims == 2) { int64_t patch_area = config.patch_size * config.patch_size; config.out_channels = tensor_storage.ne[1] / patch_area; diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index d15d754080..b94d7fdb86 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -63,16 +63,16 @@ namespace LingBotVideo { } if (ends_with(name, "patch_embedder.weight") && tensor_storage.n_dims == 2) { - int64_t patch_dim = tensor_storage.ne[0]; + int64_t patch_dim = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1]; int64_t patch_volume = config.patch_t * config.patch_h * config.patch_w; if (patch_dim % patch_volume == 0) { config.in_channels = patch_dim / patch_volume; } } else if (ends_with(name, "text_embedder.linear_1.weight") && tensor_storage.n_dims == 2) { - config.text_dim = tensor_storage.ne[0]; + config.text_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "time_embedder.linear_1.weight") && tensor_storage.n_dims == 2) { - config.freq_dim = tensor_storage.ne[0]; + config.freq_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "blocks.0.attn.norm_q.weight") && tensor_storage.n_dims == 1) { int64_t head_dim = tensor_storage.ne[0]; if (head_dim > 0) { diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index 4a9ea1e5ab..6b2b534732 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -74,13 +74,13 @@ namespace LLaDAImage { } if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { int64_t patch_area = config.patch_size * config.patch_size; - config.in_channels = tensor_storage.ne[0] / patch_area; + config.in_channels = tensor_storage.logical_ne0() / patch_area; config.hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) { - config.cap_feat_dim = tensor_storage.ne[0]; + config.cap_feat_dim = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "sigvq_embedder.1.weight") && tensor_storage.n_dims == 2) { - config.semantic_feat_dim = tensor_storage.ne[0]; + config.semantic_feat_dim = tensor_storage.logical_ne0(); } else if (ends_with(name, "layers.0.attention.to_q.weight") && tensor_storage.n_dims == 2) { detected_q_dim = tensor_storage.ne[1]; } else if (ends_with(name, "layers.0.attention.to_k.weight") && tensor_storage.n_dims == 2) { diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index 32e38012bd..f5d2b36157 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -168,7 +168,7 @@ namespace LTXV { LTXAVConfig config; auto patchify_proj_iter = tensor_storage_map.find(prefix + ".patchify_proj.weight"); if (patchify_proj_iter != tensor_storage_map.end()) { - config.in_channels = patchify_proj_iter->second.ne[0]; + config.in_channels = patchify_proj_iter->second.logical_ne0(); config.hidden_size = patchify_proj_iter->second.ne[1]; int64_t video_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.attn1.to_gate_logits.bias", 32); auto attn_layout = infer_attention_layout(config.hidden_size, video_heads); @@ -178,7 +178,7 @@ namespace LTXV { auto audio_patchify_proj_iter = tensor_storage_map.find(prefix + ".audio_patchify_proj.weight"); if (audio_patchify_proj_iter != tensor_storage_map.end()) { - config.audio_in_channels = audio_patchify_proj_iter->second.ne[0]; + config.audio_in_channels = audio_patchify_proj_iter->second.logical_ne0(); config.audio_hidden_size = audio_patchify_proj_iter->second.ne[1]; config.audio_out_channels = config.audio_in_channels; int64_t audio_heads = infer_gate_heads(tensor_storage_map, prefix + ".transformer_blocks.0.audio_attn1.to_gate_logits.bias", 32); @@ -198,11 +198,11 @@ namespace LTXV { auto attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.attn2.to_k.weight"); if (attn2_iter != tensor_storage_map.end()) { - config.cross_attention_dim = attn2_iter->second.ne[0]; + config.cross_attention_dim = attn2_iter->second.logical_ne0(); } auto audio_attn2_iter = tensor_storage_map.find(prefix + ".transformer_blocks.0.audio_attn2.to_k.weight"); if (audio_attn2_iter != tensor_storage_map.end()) { - config.audio_cross_attention_dim = audio_attn2_iter->second.ne[0]; + config.audio_cross_attention_dim = audio_attn2_iter->second.logical_ne0(); } if (tensor_storage_map.find(prefix + ".transformer_blocks.0.prompt_scale_shift_table") != tensor_storage_map.end()) { config.cross_attention_adaln = true; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index a7c2f1f930..b5aec4397d 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -73,16 +73,16 @@ namespace MiniMaxH3 { }; if (const auto* weight = find("video_patch_proj.weight")) { - config.video_latent_channels = weight->ne[0] / 4; + config.video_latent_channels = weight->logical_ne0() / 4; config.hidden_size = weight->ne[1]; } if (const auto* weight = find("audio_patch_proj.weight")) { - config.audio_latent_channels = weight->ne[0]; + config.audio_latent_channels = weight->logical_ne0(); } config.num_layers = count_blocks(tensors, prefix + ".blocks."); config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks."); if (const auto* weight = find("blocks.0.attn.q_norm.weight")) { - config.attention_head_dim = weight->ne[0]; + config.attention_head_dim = weight->logical_ne0(); } if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) { config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim); @@ -91,14 +91,14 @@ namespace MiniMaxH3 { config.ffn_hidden_size = weight->ne[1] / 2; } if (const auto* weight = find("condition_proj.weight")) { - config.text_dim = weight->ne[0]; + config.text_dim = weight->logical_ne0(); } if (const auto* table = find("adaln_t_table")) { config.time_embed_dim = table->ne[0]; config.adaln_curve_grid = table->ne[1]; } else { if (const auto* weight = find("time_embedder.proj_in.weight")) { - config.timestep_input_dim = weight->ne[0]; + config.timestep_input_dim = weight->logical_ne0(); config.time_embed_hidden_size = weight->ne[1]; } if (const auto* weight = find("time_embedder.proj_out.weight")) { diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index c6e90dae1c..5029759d33 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -59,14 +59,14 @@ namespace MiniT2I { config.pca_channels = tensor_storage.ne[2]; config.hidden_size = tensor_storage.ne[3]; } else if (ends_with(name, "txt_embedder.weight") && tensor_storage.n_dims == 2) { - config.txt_input_size = tensor_storage.ne[0]; + config.txt_input_size = tensor_storage.logical_ne0(); config.txt_hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "pooled_embedder.weight") && tensor_storage.n_dims == 2) { config.cond_vec_size = tensor_storage.ne[1]; } else if (ends_with(name, "double_blocks.0.img_qkv.weight") && tensor_storage.n_dims == 2) { int64_t inner3 = tensor_storage.ne[1]; int64_t inner = inner3 / 3; - config.hidden_size = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.logical_ne0(); if (config.hidden_size == 768) { config.num_heads = 12; config.head_dim = 64; @@ -79,7 +79,7 @@ namespace MiniT2I { } } else if (ends_with(name, "final_layer.linear.weight") && tensor_storage.n_dims == 2) { int64_t patch_area = config.patch_size * config.patch_size; - config.hidden_size = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.logical_ne0(); config.in_channels = patch_area > 0 ? tensor_storage.ne[1] / patch_area : config.in_channels; } else if (ends_with(name, "mask_token") && tensor_storage.n_dims >= 2) { config.prompt_length = tensor_storage.ne[1]; diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index 613cfc11a7..b032273db5 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -61,17 +61,17 @@ struct MMDiTConfig { } else if (name.find("y_embedder.mlp.0.weight") != std::string::npos && tensor_storage.n_dims == 2) { has_weight_config = true; has_hidden_size = true; - config.adm_in_channels = tensor_storage.ne[0]; + config.adm_in_channels = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1]; } else if (name.find("context_embedder.weight") != std::string::npos && tensor_storage.n_dims == 2) { has_weight_config = true; has_context_embed = true; - config.context_size = tensor_storage.ne[0]; + config.context_size = tensor_storage.logical_ne0(); config.context_embedder_out_dim = tensor_storage.ne[1]; } else if (name.find("final_layer.linear.weight") != std::string::npos && tensor_storage.n_dims == 2) { has_weight_config = true; has_hidden_size = true; - config.hidden_size = tensor_storage.ne[0]; + config.hidden_size = tensor_storage.logical_ne0(); int64_t patch_area = static_cast(config.patch_size) * config.patch_size; if (patch_area > 0) { config.out_channels = tensor_storage.ne[1] / patch_area; diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index a0fae7448f..2641417abf 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -32,16 +32,16 @@ namespace MOSSImage { } auto key = name.substr(prefix.size() + 1); if (key == "video_embedder.weight") { - config.channels = tensor.ne[0]; + config.channels = tensor.logical_ne0(); config.dim = tensor.ne[1]; } else if (key == "text_embedder.proj.weight") { - config.text_dim = tensor.ne[0]; + config.text_dim = tensor.logical_ne0(); } else if (key == "layers.0.proj.ffn_down.weight") { - config.ffn_dim = tensor.ne[0]; + config.ffn_dim = tensor.logical_ne0(); } else if (key == "layers.0.q_norm.weight") { config.head_dim = tensor.ne[0]; } else if (key == "time_embedding.0.weight") { - config.freq_dim = tensor.ne[0]; + config.freq_dim = tensor.logical_ne0(); } if (starts_with(key, "layers.")) { depth = std::max(depth, std::stoi(key.substr(7)) + 1); diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp index dc6e2c2e7d..74bde6d453 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -51,10 +51,10 @@ namespace PixArt { config.out_channels = w->ne[1] / (config.patch_size * config.patch_size); } if (auto w = find("y_embedder.y_proj.fc1.weight")) { - config.caption_channels = w->ne[0]; + config.caption_channels = w->logical_ne0(); } if (auto w = find("blocks.0.cross_attn.kv_linear.weight")) { - config.cross_attention_dim = w->ne[0]; + config.cross_attention_dim = w->logical_ne0(); } if (auto w = find("blocks.0.mlp.fc1.weight")) { config.ffn_dim = w->ne[1]; diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 76cbccb838..9452ec4678 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -24,14 +24,14 @@ namespace Qwen { return it == weights.end() ? nullptr : &it->second; }; if (auto w = find("img_in.weight")) { - config.in_channels = w->ne[0]; + config.in_channels = w->logical_ne0(); config.hidden_size = w->ne[1]; } if (auto w = find("proj_out.weight")) { config.out_channels = w->ne[1]; } if (auto w = find("txt_in.in_layer.weight")) { - config.context_dim = w->ne[0]; + config.context_dim = w->logical_ne0(); } if (auto w = find("transformer_blocks.0.attn.norm_q.weight")) { config.head_dim = w->ne[0]; diff --git a/src/model/diffusion/sefi_image.hpp b/src/model/diffusion/sefi_image.hpp index 6782b68b49..7781532260 100644 --- a/src/model/diffusion/sefi_image.hpp +++ b/src/model/diffusion/sefi_image.hpp @@ -31,7 +31,7 @@ namespace SefiImage { continue; } if (ends_with(name, "dual_time_embed.semantic_embedder.linear_1.weight") && tensor_storage.n_dims == 2) { - config.timestep_guidance_in_dim = tensor_storage.ne[0]; + config.timestep_guidance_in_dim = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1] * 2; } } diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp index 7c0ff09506..675d94296c 100644 --- a/src/model/diffusion/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -104,13 +104,13 @@ struct UNetConfig { } if (const TensorStorage* time_embed = find_weight("time_embed.0.weight")) { if (time_embed->n_dims == 2) { - config.model_channels = static_cast(time_embed->ne[0]); + config.model_channels = static_cast(time_embed->logical_ne0()); config.time_embed_dim = static_cast(time_embed->ne[1]); } } if (const TensorStorage* label_emb = find_weight("label_emb.0.0.weight")) { if (label_emb->n_dims == 2) { - config.adm_in_channels = static_cast(label_emb->ne[0]); + config.adm_in_channels = static_cast(label_emb->logical_ne0()); config.time_embed_dim = static_cast(label_emb->ne[1]); } } @@ -124,7 +124,7 @@ struct UNetConfig { continue; } if (name.find("attn2.to_k.weight") != std::string::npos && tensor_storage.n_dims == 2) { - config.context_dim = static_cast(tensor_storage.ne[0]); + config.context_dim = static_cast(tensor_storage.logical_ne0()); break; } } diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index 1c66c96b7f..bc9ddc4498 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -62,10 +62,10 @@ namespace ZImage { } if (ends_with(name, "x_embedder.weight") && tensor_storage.n_dims == 2) { int64_t patch_area = config.patch_size * config.patch_size; - config.in_channels = tensor_storage.ne[0] / patch_area; + config.in_channels = tensor_storage.logical_ne0() / patch_area; config.hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "cap_embedder.1.weight") && tensor_storage.n_dims == 2) { - config.cap_feat_dim = tensor_storage.ne[0]; + config.cap_feat_dim = tensor_storage.logical_ne0(); config.hidden_size = tensor_storage.ne[1]; } else if (ends_with(name, "layers.0.attention.q_norm.weight") && tensor_storage.n_dims == 1) { detected_head_dim = tensor_storage.ne[0]; diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp index f44f2ff9f8..2ba6680329 100644 --- a/src/model/diffusion/z_image_l2p.hpp +++ b/src/model/diffusion/z_image_l2p.hpp @@ -18,7 +18,7 @@ namespace ZImageL2P { config.patch_size = 16; auto x_embedder = tensors.find(prefix + ".x_embedder.weight"); if (x_embedder != tensors.end()) { - config.patch_size = static_cast(std::lround(std::sqrt(static_cast(x_embedder->second.ne[0] / config.in_channels)))); + config.patch_size = static_cast(std::lround(std::sqrt(static_cast(x_embedder->second.logical_ne0() / config.in_channels)))); } // L2P ships split to_q/to_k/to_v, which name conversion maps to qkv.weight plus // .1/.2 parts; the base detection only sees the q part and undercounts kv heads. diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index 7e7ac30984..29c6f87311 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -57,7 +57,7 @@ struct T5Config { } if (const TensorStorage* q = find_tensor("encoder.block.0.layer.0.SelfAttention.q.weight")) { if (q->n_dims == 2) { - config.model_dim = q->ne[0]; + config.model_dim = q->logical_ne0(); int64_t inner_dim = q->ne[1]; config.inner_dim = inner_dim; // Flan-T5/T5 uses d_kv=64 for common sizes. @@ -68,7 +68,7 @@ struct T5Config { } if (const TensorStorage* wi = find_tensor("encoder.block.0.layer.1.DenseReluDense.wi_0.weight")) { if (wi->n_dims == 2) { - config.model_dim = wi->ne[0]; + config.model_dim = wi->logical_ne0(); config.ff_dim = wi->ne[1]; } } diff --git a/src/model/upscaler/h3_latent_upscaler.hpp b/src/model/upscaler/h3_latent_upscaler.hpp index 006415081f..80b96e7395 100644 --- a/src/model/upscaler/h3_latent_upscaler.hpp +++ b/src/model/upscaler/h3_latent_upscaler.hpp @@ -75,7 +75,7 @@ namespace H3VUpsampler { 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]); + config.embed_dim = static_cast(embed_mid->logical_ne0()); } std::set in_res_idx, in_temporal_idx, out_res_idx, out_temporal_idx; diff --git a/src/model_io/safetensors_io.cpp b/src/model_io/safetensors_io.cpp index cf424af921..058606267d 100644 --- a/src/model_io/safetensors_io.cpp +++ b/src/model_io/safetensors_io.cpp @@ -109,8 +109,9 @@ static ggml_type safetensors_dtype_to_ggml_type(const std::string& dtype) { struct ComfyQuantConfig { std::string format; - bool convrot = false; - int group_size = 0; + bool convrot = false; + int group_size = 0; + int weight_group_size = 0; }; static bool read_comfy_quant_config(std::ifstream& file, @@ -140,6 +141,12 @@ static bool read_comfy_quant_config(std::ifstream& file, config.format = json.value("format", ""); config.convrot = json.value("convrot", false); config.group_size = json.value("convrot_groupsize", 0); + config.weight_group_size = json.value("group_size", 0); + if (config.format == "convrot_w4a4" || config.format == "asym_w4a8_int8") { + // Packed W4 formats imply rotation, including exports without the INT8 convrot flag. + config.convrot = true; + config.group_size = json.value("convrot_groupsize", 256); + } } catch (const std::exception&) { set_error(error, "parsing ComfyUI quantization metadata tensor failed: '" + tensor_name + "'"); return false; @@ -238,8 +245,8 @@ bool read_safetensors_file(const std::string& file_path, return false; } #ifdef SD_USE_UPSTREAM_GGML - if (config.format == "int8_tensorwise") { - set_error(error, "INT8 tensorwise/convrot is not supported by this ggml build (tensor '" + name + "')"); + if (config.format == "int8_tensorwise" || config.format == "convrot_w4a4" || config.format == "asym_w4a8_int8") { + set_error(error, "ComfyUI INT8/packed INT4 is not supported by this ggml build (tensor '" + name + "')"); return false; } #endif @@ -340,29 +347,58 @@ bool read_safetensors_file(const std::string& file_path, if (ends_with(name, ".weight")) { const std::string module_name = name.substr(0, name.size() - std::string(".weight").size()); auto config = comfy_quant_configs.find(module_name); - if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise") { + if (config != comfy_quant_configs.end() && + (config->second.format == "int8_tensorwise" || config->second.format == "convrot_w4a4" || config->second.format == "asym_w4a8_int8")) { + const int w4_kind = config->second.format == "convrot_w4a4" ? 1 : config->second.format == "asym_w4a8_int8" ? 2 + : 0; if (type != GGML_TYPE_I8) { - set_error(error, "ComfyUI int8_tensorwise weight is not I8: '" + name + "'"); + set_error(error, "ComfyUI quantized weight is not I8: '" + name + "'"); + return false; + } + if (w4_kind && (tensor_storage.n_dims != 2 || tensor_storage.ne[0] <= 0 || + tensor_storage.ne[0] > INT64_MAX / 2 || tensor_storage.ne[1] <= 0)) { + set_error(error, "invalid packed INT4 weight shape: '" + name + "'"); return false; } + const int64_t input_features = tensor_storage.ne[0] * (w4_kind ? 2 : 1); if (config->second.convrot) { int group_size_remainder = config->second.group_size; while (group_size_remainder > 1 && group_size_remainder % 4 == 0) { group_size_remainder /= 4; } - if (group_size_remainder != 1 || tensor_storage.ne[0] % config->second.group_size != 0) { + if (group_size_remainder != 1 || config->second.group_size > 256 || input_features % config->second.group_size != 0) { set_error(error, "invalid ComfyUI convrot group size for tensor '" + name + "'"); return false; } } - tensor_storage.is_int8_tensorwise = true; + if (w4_kind) { + auto matches = [&](const std::string& suffix, const std::string& expected_dtype, const std::vector& expected_shape) { + auto entry = header_.find(module_name + suffix); + return entry != header_.end() && entry->value("dtype", "") == expected_dtype && + (*entry)["shape"] == expected_shape; + }; + const int64_t n = tensor_storage.ne[1]; + const bool valid = w4_kind == 1 + ? (matches(".weight_scale", "F32", {n}) || matches(".weight_scale", "F32", {n, 1})) + : config->second.weight_group_size == 16 && input_features % 16 == 0 && + matches(".weight_codebook", "F32", {16}) && + (matches(".weight_s_channel", "F32", {n}) || matches(".weight_s_channel", "F32", {n, 1})) && + matches(".weight_s_rel", "F8_E4M3", {n, input_features / 16}); + if (!valid) { + set_error(error, "invalid packed INT4 scales/codebook for tensor '" + name + "'"); + return false; + } + } + tensor_storage.is_int8_tensorwise = w4_kind == 0; + tensor_storage.w4_convrot_kind = w4_kind; tensor_storage.int8_convrot = config->second.convrot; tensor_storage.int8_convrot_group_size = config->second.group_size; } - } else if (ends_with(name, ".weight_scale")) { - const std::string module_name = name.substr(0, name.size() - std::string(".weight_scale").size()); + } else if (ends_with(name, ".weight_scale") || ends_with(name, ".weight_s_channel")) { + const std::string module_name = name.substr(0, name.rfind('.')); auto config = comfy_quant_configs.find(module_name); - if (config != comfy_quant_configs.end() && config->second.format == "int8_tensorwise" && + if (config != comfy_quant_configs.end() && + (config->second.format == "int8_tensorwise" || config->second.format == "convrot_w4a4" || config->second.format == "asym_w4a8_int8") && tensor_storage.n_dims == 2 && tensor_storage.ne[0] == 1) { tensor_storage.ne[0] = tensor_storage.ne[1]; tensor_storage.ne[1] = 1; diff --git a/src/model_io/tensor_storage.h b/src/model_io/tensor_storage.h index 1916d1f5e5..1062174ebc 100644 --- a/src/model_io/tensor_storage.h +++ b/src/model_io/tensor_storage.h @@ -22,6 +22,7 @@ struct TensorStorage { bool is_f64 = false; bool is_i64 = false; bool is_int8_tensorwise = false; + int w4_convrot_kind = 0; bool int8_convrot = false; int int8_convrot_group_size = 0; int64_t ne[SD_MAX_DIMS] = {1, 1, 1, 1, 1}; @@ -51,6 +52,10 @@ struct TensorStorage { return n; } + int64_t logical_ne0() const { + return ne[0] * (w4_convrot_kind ? 2 : 1); + } + int64_t nbytes() const { return nelements() * ggml_type_size(type) / ggml_blck_size(type); } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 407885bbff..73025a4075 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -518,7 +518,7 @@ SDVersion ModelLoader::get_sd_version() const { } if (tensor_storage.name.find("model.diffusion_model.transformer_blocks.0.img_mod.1.weight") != std::string::npos) { auto img_in = tensor_storage_map.find("model.diffusion_model.img_in.weight"); - if (img_in != tensor_storage_map.end() && img_in->second.ne[0] == 128) { + if (img_in != tensor_storage_map.end() && img_in->second.logical_ne0() == 128) { return VERSION_MAGE_FLOW; } if (tensor_storage_map.find("model.diffusion_model.time_text_embed.addition_t_embedding.weight") != tensor_storage_map.end()) { @@ -593,7 +593,7 @@ SDVersion ModelLoader::get_sd_version() const { } } auto connector = tensor_storage_map.find("embeddings_connectors.proj.weight"); - if (connector != tensor_storage_map.end() && connector->second.ne[0] == 12800 && + if (connector != tensor_storage_map.end() && connector->second.logical_ne0() == 12800 && connector->second.ne[1] == 5120 && tensor_storage_map.count("embeddings_connectors.sink_embedding")) { return VERSION_P_H3_IMAGE; } @@ -660,7 +660,7 @@ SDVersion ModelLoader::get_sd_version() const { 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) { - if (tensor_storage.ne[0] == 1024) + if (tensor_storage.logical_ne0() == 1024) has_attn_1024 = true; } } @@ -721,16 +721,16 @@ SDVersion ModelLoader::get_sd_version() const { if (has_distilled_guidance) { return VERSION_CHROMA; } - if (context_ebedding_weight.ne[0] == 3584) { + if (context_ebedding_weight.logical_ne0() == 3584) { return VERSION_LONGCAT; } else { - if (input_block_weight.ne[0] == 384) { + if (input_block_weight.logical_ne0() == 384) { return VERSION_FLUX_FILL; } - if (input_block_weight.ne[0] == 128) { + if (input_block_weight.logical_ne0() == 128) { return VERSION_FLUX_CONTROLS; } - if (input_block_weight.ne[0] == 196) { + if (input_block_weight.logical_ne0() == 196) { return VERSION_FLEX_2; } return VERSION_FLUX; @@ -1651,7 +1651,8 @@ bool ModelLoader::load_tensors(std::map& tensors, bool ModelLoader::tensor_should_be_converted(const TensorStorage& tensor_storage, ggml_type type) const { const std::string& name = tensor_storage.name; - if (tensor_storage.is_int8_tensorwise) { + if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind != 0 || + ends_with(name, ".weight_codebook") || ends_with(name, ".weight_s_channel") || ends_with(name, ".weight_s_rel")) { return false; } if (type != GGML_TYPE_COUNT) { diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp index 10c7d2e68f..adeea95d22 100644 --- a/src/model_manager_files.cpp +++ b/src/model_manager_files.cpp @@ -11,7 +11,7 @@ static bool same_tensor_source(const TensorStorage& a, const TensorStorage& b) { a.n_dims == b.n_dims && std::equal(a.ne, a.ne + SD_MAX_DIMS, b.ne) && a.is_f8_e4m3 == b.is_f8_e4m3 && a.is_f8_e5m2 == b.is_f8_e5m2 && a.is_f64 == b.is_f64 && a.is_i64 == b.is_i64 && - a.is_int8_tensorwise == b.is_int8_tensorwise && a.int8_convrot == b.int8_convrot && + a.is_int8_tensorwise == b.is_int8_tensorwise && a.w4_convrot_kind == b.w4_convrot_kind && a.int8_convrot == b.int8_convrot && a.int8_convrot_group_size == b.int8_convrot_group_size; } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 98e80e1174..4875858b82 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -1019,7 +1019,7 @@ void StableDiffusionGGML::configure_weight_loading() { bool have_int8_tensorwise = false; bool have_quant_type_rule = false; for (const auto& [_, tensor_storage] : model_loader.get_tensor_storage_map()) { - if (tensor_storage.is_int8_tensorwise) { + if (tensor_storage.is_int8_tensorwise || tensor_storage.w4_convrot_kind) { have_int8_tensorwise = true; } if (tensor_storage.expected_type < GGML_TYPE_COUNT && ggml_is_quantized(tensor_storage.expected_type)) { @@ -1051,7 +1051,7 @@ void StableDiffusionGGML::configure_weight_loading() { } else if (sd_ctx_params->lora_apply_mode == LORA_APPLY_IMMEDIATELY) { if (have_int8_tensorwise) { LOG_WARN( - "INT8 tensorwise weights do not support the immediately LoRA apply mode; " + "ComfyUI INT8/packed INT4 weights do not support the immediately LoRA apply mode; " "using at_runtime instead"); apply_lora_immediately = false; } else if (row_split_active()) { From f5df5f8b84f139f1e80496902da90798a0d92896 Mon Sep 17 00:00:00 2001 From: leejet Date: Fri, 6 Nov 2026 01:59:05 +0800 Subject: [PATCH 21/57] perf: accelerate packed INT4 CUDA inference with CUTLASS (#2278) --- docs/int8_convrot.md | 19 +++-- src/core/cuda/CMakeLists.txt | 6 ++ src/core/cuda/i8_gemm.cu | 32 ++++---- src/core/cuda/i8_gemm.h | 2 +- src/core/cuda/w4.cu | 76 +++++++++++++++++-- src/core/cuda/w4_mma.h | 138 +++++++++++++++++++++++++++++++++++ 6 files changed, 247 insertions(+), 26 deletions(-) create mode 100644 src/core/cuda/w4_mma.h diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index 49e77bcefc..f1e921aa7e 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -123,13 +123,18 @@ Both formats imply convrot rotation even without a `convrot` flag; missing Both formats reuse the convrot activation cache when the input scale is `1` and produce F32 outputs with optional bias. CPU provides a portable implementation. -CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS; -CUDA and supported AMD targets use packed byte-dot instructions, with portable -byte arithmetic for other targets. Vulkan requires integer dot product support -and at least 13312 bytes of shared memory for W4A4 or 16452 bytes for W4A8. -GPU W4 kernels require the logical input width to be divisible by 64; H64 and -H256 are supported. Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA -by the active compilation target's +CUDA/HIP/MUSA use shared packed-weight kernels without requiring INT8 BLAS. +With CUTLASS enabled, NVIDIA +SM80 and newer devices use mixed-input S8 x S4 GEMM for eligible W4A4 shapes, +reusing the INT8 scale epilogue without expanding the full weight tensor. W4A8 +uses CUTLASS integer MMA per 16-value group with FP32 scale restoration. +Other CUDA shapes and builds retain the packed byte-dot path. Set +`GGML_CUDA_W4_FORCE_DP4A=1` to compare against that path. Supported AMD targets +use packed byte-dot instructions, with portable byte arithmetic for other targets. +Vulkan requires integer dot product support and at least 13312 bytes of shared +memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input +width to be divisible by 64; H64 and H256 are supported. Other shapes use CPU +fallback. HIP W4A8 selects wave64 WMMA by the active compilation target's intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout. There is no GPU architecture allowlist: newer targets supporting either intrinsic automatically use its matching layout. Selection happens per device diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt index 593445bf0e..e2dd3f92c3 100644 --- a/src/core/cuda/CMakeLists.txt +++ b/src/core/cuda/CMakeLists.txt @@ -27,6 +27,12 @@ if(SD_CUDA) set_source_files_properties(sage_attention.cu sol_attention.cu PROPERTIES COMPILE_OPTIONS "--use_fast_math") set_source_files_properties(${SD_CUDA_SOURCES} PROPERTIES COMPILE_OPTIONS "--use_fast_math") if(SD_CUDA_CUTLASS_PATH) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_W4_CUTLASS) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY INCLUDE_DIRECTORIES "${SD_CUDA_CUTLASS_PATH}/include") + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS --expt-relaxed-constexpr) + if(MSVC) + set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/w4_cuda.cu" APPEND PROPERTY COMPILE_OPTIONS -Xcompiler=/Zc:preprocessor,/Zc:__cplusplus) + endif() set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/bf16_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_BF16_CUTLASS) set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/int8_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_I8_CUTLASS) set_property(SOURCE "${CMAKE_CURRENT_BINARY_DIR}/registration_cuda.cu" APPEND PROPERTY COMPILE_DEFINITIONS SD_CUDA_CUTE_ATTENTION) diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu index 69f91d2aea..9d4d92e1ca 100644 --- a/src/core/cuda/i8_gemm.cu +++ b/src/core/cuda/i8_gemm.cu @@ -143,8 +143,9 @@ namespace sd_i8_cutlass { dst[i] = src[i / 2 + (i % 2) * (n / 2)]; } - template + template struct DenseTypes { + using Weight = std::conditional_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; using Swizzle = std::conditional_t>; @@ -161,10 +162,10 @@ namespace sd_i8_cutlass { cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, - int8_t, + Weight, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, - 16, + PackedW4 ? 32 : 16, float, cutlass::layout::RowMajor, 4, @@ -178,11 +179,12 @@ namespace sd_i8_cutlass { Epilogue, Swizzle, 4, - cutlass::arch::OpMultiplyAddSaturate>::GemmKernel; + std::conditional_t>::GemmKernel; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; template struct SwigluTypes { + using Weight = int8_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; using Swizzle = std::conditional_t>; @@ -205,6 +207,7 @@ namespace sd_i8_cutlass { DeviceInfo device; int64_t rows, n, k; bool swiglu; + bool packed_w4; 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; @@ -213,7 +216,7 @@ namespace sd_i8_cutlass { 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}; + return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast(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) { @@ -272,6 +275,8 @@ namespace sd_i8_cutlass { } template static bool configure_tile(Plan& p) { + if (p.packed_w4) + return configure_kernel>(p); if (p.swiglu) return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); @@ -300,20 +305,21 @@ namespace sd_i8_cutlass { } } 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) { +SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) { #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) || + if ((packed_w4 && (swiglu || k > INT_MAX / 1024)) || info.major < 8 || rows < (sd_i8_cutlass::small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n < 1 || n > INT_MAX || k < 1 || k > INT_MAX || k % 64 || n % (swiglu ? 256 : 128)) return nullptr; - auto* p = new SDI8CudaGemmPlan; - p->device = info; - p->rows = rows; - p->n = n; - p->k = k; - p->swiglu = swiglu; + auto* p = new SDI8CudaGemmPlan; + p->device = info; + p->rows = rows; + p->n = n; + p->k = k; + p->swiglu = swiglu; + p->packed_w4 = packed_w4; if (!sd_i8_cutlass::configure(*p)) { delete p; return nullptr; diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h index 189df6ef0b..5e4ac9f8c2 100644 --- a/src/core/cuda/i8_gemm.h +++ b/src/core/cuda/i8_gemm.h @@ -3,7 +3,7 @@ #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); +SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false); ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream); void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan); #endif diff --git a/src/core/cuda/w4.cu b/src/core/cuda/w4.cu index ef46d02369..98f810e136 100644 --- a/src/core/cuda/w4.cu +++ b/src/core/cuda/w4.cu @@ -1,6 +1,13 @@ #include #include "core/ggml_extensions.h" #include "runtime.h" +#ifdef SD_CUDA_W4_CUTLASS +#include +#include +#include +#include "i8_gemm.h" +#include "w4_mma.h" +#endif namespace SD_CUDA_NAMESPACE { namespace sd_w4_cuda { @@ -789,7 +796,37 @@ namespace SD_CUDA_NAMESPACE { return false; return true; } - static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) { + struct Plan { + bool mma = false; + size_t gemm_offset = 0; +#ifdef SD_CUDA_W4_CUTLASS + SDI8CudaGemmPlan* gemm = nullptr; + ~Plan() { sd_cuda_i8_gemm_free(gemm); } +#endif + }; + static bool use_mma(ggml_backend_t backend, const SDW4Operation& op) { +#ifdef SD_CUDA_W4_CUTLASS + if (!aligned(op.input, 16) || !aligned(op.weight, 16)) + return false; + static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_W4_FORCE_DP4A"); return v && std::atoi(v); }(); + if (disabled) + return false; + static std::mutex mutex; + static std::unordered_map devices; + const auto device = cuda_device_info(ggml_backend_get_device(backend)).device; + std::lock_guard lock(mutex); + auto found = devices.find(device); + if (found != devices.end()) + return found->second; + cudaDeviceProp prop = {}; + const bool available = cudaGetDeviceProperties(&prop, device) == cudaSuccess && prop.major >= 8; + devices.emplace(device, available); + return available; +#else + return false; +#endif + } + static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int, void** state, size_t* bytes, void*) { SDW4Operation op; if (!sd_ggml_w4_describe(nodes[0], &op)) return GGML_STATUS_FAILED; @@ -812,7 +849,22 @@ namespace SD_CUDA_NAMESPACE { may_alias |= !input->buffer || (input->buffer == dst->buffer && a < b + size && b < a + ggml_nbytes(dst)); } } - *bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 ? 1 : 0) + (may_alias ? 1 : 0)); + auto* plan = new Plan; + plan->mma = use_mma(backend, op); + *state = plan; + size_t gemm_bytes = 0; +#ifdef SD_CUDA_W4_CUTLASS + if (plan->mma && op.kind == SDW4Kind::W4A4) { + if (!op.bias && aligned(dst, 16) && aligned(op.weight_scale, 16)) + plan->gemm = sd_cuda_i8_gemm_prepare(ggml_backend_get_device(backend), ggml_nrows(dst), dst->ne[0], op.input->ne[0], false, &gemm_bytes, true); + plan->mma = plan->gemm != nullptr; + if (plan->gemm && gemm_bytes == 0) + gemm_bytes = 256; + } +#endif + const size_t tensor_bytes = ggml_nbytes(dst) * ((op.kind == SDW4Kind::W4A4 && !plan->mma ? 1 : 0) + (may_alias ? 1 : 0)); + plan->gemm_offset = (tensor_bytes + 255) / 256 * 256; + *bytes = plan->gemm_offset + gemm_bytes; return GGML_STATUS_SUCCESS; } static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) { @@ -821,7 +873,7 @@ namespace SD_CUDA_NAMESPACE { 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 ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void*, void*) { + static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int, void* state, void*) { auto* dst = nodes[0]; SDW4Operation op; ggml_ext_gpu_context context; @@ -840,7 +892,21 @@ namespace SD_CUDA_NAMESPACE { const auto* ws = static_cast(op.weight_scale->data); const auto* bias = op.bias ? static_cast(op.bias->data) : nullptr; const dim3 grid(unsigned((n + 63) / 64), unsigned((rows + 31) / 32)), block(16, 16); - if (op.kind == SDW4Kind::W4A4) { +#ifdef SD_CUDA_W4_CUTLASS + auto* plan = static_cast(state); + if (plan->mma) { + if (plan->gemm) { + auto* workspace = scratch ? reinterpret_cast(scratch) + plan->gemm_offset : nullptr; + if (sd_cuda_i8_gemm_execute(plan->gemm, reinterpret_cast(x), reinterpret_cast(w), xs, ws, output, workspace, stream) != GGML_STATUS_SUCCESS) + return GGML_STATUS_FAILED; + } else { + const dim3 tiles(unsigned((n + 63) / 64), unsigned((rows + 127) / 128)); + sd_w4_mma::w4a8<128, 64, 128><<>>(reinterpret_cast(w), reinterpret_cast(x), xs, ws, + static_cast(op.s_channel->data), static_cast(op.s_rel->data), bias, output, int(rows), int(n), int(k)); + } + } else +#endif + if (op.kind == SDW4Kind::W4A4) { auto* accum = reinterpret_cast(scratch + (alias ? ggml_nelements(dst) : 0)); mul_mat_w4a4_convrot_dp4a<<>>(w, x, accum, uint32_t(n), uint32_t(rows), uint32_t(k / 2)); if (cuda_last_error() != cuda_success) @@ -867,7 +933,7 @@ namespace SD_CUDA_NAMESPACE { d.supports = sd_w4_cuda::supports; d.prepare = sd_w4_cuda::prepare; d.execute = sd_w4_cuda::execute; - d.destroy = [](void*, void*) {}; + d.destroy = [](void* state, void*) { delete static_cast(state); }; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &d, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } } diff --git a/src/core/cuda/w4_mma.h b/src/core/cuda/w4_mma.h new file mode 100644 index 0000000000..94629bfbd6 --- /dev/null +++ b/src/core/cuda/w4_mma.h @@ -0,0 +1,138 @@ +#ifndef __SD_CORE_CUDA_W4_MMA_H__ +#define __SD_CORE_CUDA_W4_MMA_H__ + +#include +#include +#include + +namespace sd_w4_mma { + static __device__ __forceinline__ float fp8_e4m3(uint8_t v) { + const int exponent = (v >> 3) & 15, mantissa = v & 7; + const float result = exponent == 15 && mantissa == 7 ? __int_as_float(0x7fc00000) : exponent == 0 ? ldexpf(float(mantissa), -9) + : ldexpf(1.f + float(mantissa) / 8.f, exponent - 7); + return v & 128 ? -result : result; + } + + template + static __global__ void __launch_bounds__(256) w4a8( + const uint8_t* __restrict__ w, + const int8_t* __restrict__ x, + const float* __restrict__ xs, + const float* __restrict__ codebook, + const float* __restrict__ channel, + const uint8_t* __restrict__ relative, + const float* __restrict__ bias, + float* __restrict__ dst, + int rows, + int n, + int k) { + const int row_base = blockIdx.y * M, col_base = blockIdx.x * N; + const int tid = threadIdx.x, lane = tid % 32, warp = tid / 32; + __shared__ int cb[16]; + __shared__ float cbscale; + if (tid == 0) { + float amax = 0.f; +#pragma unroll + for (int i = 0; i < 16; ++i) + amax = fmaxf(amax, fabsf(codebook[i])); + cbscale = amax / 127.f; +#pragma unroll + for (int i = 0; i < 16; ++i) { + const int value = cbscale == 0.f ? 0 : __float2int_rn(codebook[i] / cbscale); + cb[i] = max(-127, min(127, value)); + } + } + __syncthreads(); +#if __CUDA_ARCH__ >= 800 + constexpr int stride = K / 4 + 4; + __shared__ __align__(16) uint32_t a[M][stride], b[N][stride]; + __shared__ uint32_t lut[256]; + __shared__ float scales[K / 16][N]; + lut[tid] = (cb[tid & 15] & 255) | ((cb[tid >> 4] & 255) << 8); + __syncthreads(); + constexpr int warp_cols = N / 32; + constexpr int row_tiles = M / (16 * (8 / warp_cols)); + const int wr = (warp / warp_cols) * 16 * row_tiles, wc = (warp % warp_cols) * 32; + float acc[row_tiles][4][4] = {}; + for (int base = 0; base < k; base += K) { + for (int i = tid; i < M * K / 16; i += 256) { + const int r = i / (K / 16), t = (i % (K / 16)) * 16; + const int4 value = row_base + r < rows && base + t < k ? *reinterpret_cast(x + int64_t(row_base + r) * k + base + t) : make_int4(0, 0, 0, 0); + *reinterpret_cast(&a[r][t / 4]) = value; + } + for (int i = tid; i < N * K / 8; i += 256) { + const int c = i / (K / 8), t = i % (K / 8); + const uint32_t value = col_base + c < n && base + t * 8 < k ? *reinterpret_cast(w + int64_t(col_base + c) * (k / 2) + base / 2 + t * 4) : 0; + b[c][2 * t] = lut[value & 255] | (lut[(value >> 8) & 255] << 16); + b[c][2 * t + 1] = lut[(value >> 16) & 255] | (lut[value >> 24] << 16); + } + for (int i = tid; i < N * K / 16; i += 256) { + const int c = i % N, g = i / N; + scales[g][c] = col_base + c < n && base + g * 16 < k ? cbscale * (channel[col_base + c] * fp8_e4m3(relative[int64_t(col_base + c) * (k / 16) + base / 16 + g])) : 0.f; + } + __syncthreads(); +#pragma unroll + for (int g = 0; g < K / 16; ++g) { +#pragma unroll + for (int r = 0; r < row_tiles; ++r) { + const int ar = wr + r * 16 + lane / 4, ac = g * 4 + lane % 4; + const uint32_t a0 = a[ar][ac], a1 = a[ar + 8][ac]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t b0 = b[wc + j * 8 + lane / 4][ac]; + using Mma = cutlass::arch::Mma, 32, + int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor, + int, cutlass::layout::RowMajor, cutlass::arch::OpMultiplyAddSaturate>; + Mma::FragmentA af; + Mma::FragmentB bf; + Mma::FragmentC d; + reinterpret_cast(&af)[0] = a0; + reinterpret_cast(&af)[1] = a1; + reinterpret_cast(&bf)[0] = b0; + d.clear(); + // Each K=16 MMA has its own scale; combining groups before + // dequantization would change the W4A8 representation. + Mma{}(d, af, bf, d); + const float s0 = scales[g][wc + j * 8 + (lane % 4) * 2], s1 = scales[g][wc + j * 8 + (lane % 4) * 2 + 1]; +#pragma unroll + for (int v = 0; v < 4; ++v) + acc[r][j][v] = fmaf(float(d[v]), v % 2 ? s1 : s0, acc[r][j][v]); + } + } + } + __syncthreads(); + } +#pragma unroll + for (int r = 0; r < row_tiles; ++r) { +#pragma unroll + for (int j = 0; j < 4; ++j) { +#pragma unroll + for (int v = 0; v < 4; ++v) { + const int row = row_base + wr + r * 16 + lane / 4 + (v / 2) * 8; + const int col = col_base + wc + j * 8 + (lane % 4) * 2 + v % 2; + if (row < rows && col < n) + dst[int64_t(row) * n + col] = __fmul_rn(acc[r][j][v], xs[row]) + (bias ? bias[col] : 0.f); + } + } + } +#else + for (int i = tid; i < M * N; i += 256) { + const int row = row_base + i / N, col = col_base + i % N; + if (row < rows && col < n) { + float sum = 0.f; + for (int g = 0; g < k / 16; ++g) { + int d = 0; + for (int j = 0; j < 16; ++j) { + const int t = g * 16 + j; + d += cb[(w[int64_t(col) * (k / 2) + t / 2] >> (4 * (t % 2))) & 15] * int(x[int64_t(row) * k + t]); + } + sum = fmaf(float(d), cbscale * (channel[col] * fp8_e4m3(relative[int64_t(col) * (k / 16) + g])), sum); + } + dst[int64_t(row) * n + col] = __fmul_rn(sum, xs[row]) + (bias ? bias[col] : 0.f); + } + } +#endif + } + +} +#endif From 5a1ba5d895bec4e07deab80d6d791721cd59d9c2 Mon Sep 17 00:00:00 2001 From: leejet Date: Sat, 7 Nov 2026 00:58:26 +0800 Subject: [PATCH 22/57] perf(vulkan): accelerate packed INT4 with cooperative matrices (#2280) --- docs/int8_convrot.md | 23 ++- src/core/vulkan/CMakeLists.txt | 3 + src/core/vulkan/int8.cpp | 56 +++++- .../shaders/mul_mat_w4a4_convrot_cm2.comp | 88 ++++++++++ .../mul_mat_w4a4_convrot_decode_cm2.comp | 92 ++++++++++ .../shaders/mul_mat_w4a8_convrot_cm2.comp | 164 ++++++++++++++++++ 6 files changed, 420 insertions(+), 6 deletions(-) create mode 100644 src/core/vulkan/shaders/mul_mat_w4a4_convrot_cm2.comp create mode 100644 src/core/vulkan/shaders/mul_mat_w4a4_convrot_decode_cm2.comp create mode 100644 src/core/vulkan/shaders/mul_mat_w4a8_convrot_cm2.comp diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index f1e921aa7e..d7948f6f39 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -133,9 +133,26 @@ Other CUDA shapes and builds retain the packed byte-dot path. Set use packed byte-dot instructions, with portable byte arithmetic for other targets. Vulkan requires integer dot product support and at least 13312 bytes of shared memory for W4A4 or 16452 bytes for W4A8. GPU W4 kernels require the logical input -width to be divisible by 64; H64 and H256 are supported. Other shapes use CPU -fallback. HIP W4A8 selects wave64 WMMA by the active compilation target's -intrinsic support, preferring the GFX12 fragment layout over the GFX11 layout. +width to be divisible by 64; H64 and H256 are supported. When supported by the +compiler and device, Vulkan uses NV cooperative matrix integer GEMM. W4A4 +decodes packed weights directly into cooperative matrices for aligned input +and output widths. Complete activation tiles and the final partial tile use +separate shader specializations. Other W4A4 shapes use double-buffered weight +and activation tiles (20480 bytes of shared memory, or 30720 for larger tiles). +W4A8 stages both weights and activations in shared memory in K=64 tiles, +reusing them for each 16-value group's integer +MMA and FP32 scale restoration. FP8 scales use a shared lookup table, and +expanding projections visit activation tiles in groups of weight panels for +cache reuse. Adjacent groups share a K=32 activation fragment, with +complementary zero padding in the weights to preserve each group's scale. +This path uses 17476 bytes of shared memory, or 22596 for larger tiles. +Both paths use 256-thread workgroups and keep the model's +weights packed; neither expands a full layer in device memory. Other Vulkan +devices retain the packed dot-product path; set `GGML_VK_W4_FORCE_DP4A=1` to compare +against it. FP32 accumulation order can differ between these paths. +Other shapes use CPU fallback. HIP W4A8 selects wave64 WMMA by the active +compilation target's intrinsic support, preferring the GFX12 fragment layout +over the GFX11 layout. There is no GPU architecture allowlist: newer targets supporting either intrinsic automatically use its matching layout. Selection happens per device image, including multi-architecture builds; targets supporting neither intrinsic diff --git a/src/core/vulkan/CMakeLists.txt b/src/core/vulkan/CMakeLists.txt index e71bcd5339..19367d4340 100644 --- a/src/core/vulkan/CMakeLists.txt +++ b/src/core/vulkan/CMakeLists.txt @@ -56,6 +56,9 @@ sd_vulkan_shader(quantize_i8_convrot quantize_i8_convrot ${SD_VK_DOT} -O) sd_vulkan_shader(quantize_i8_convrot_h64 quantize_i8_convrot_h64 ${SD_VK_DOT} -O -DGROUP_SIZE=64) sd_vulkan_shader(mul_mat_w4a4_convrot mul_mat_w4a4_convrot ${SD_VK_DOT} -O) sd_vulkan_shader(mul_mat_w4a8_convrot mul_mat_w4a8_convrot ${SD_VK_DOT} -O) +sd_vulkan_shader(mul_mat_w4a4_convrot_cm2 mul_mat_w4a4_convrot_cm2 ${SD_VK_CM2} -O) +sd_vulkan_shader(mul_mat_w4a4_convrot_decode_cm2 mul_mat_w4a4_convrot_decode_cm2 ${SD_VK_CM2} -O) +sd_vulkan_shader(mul_mat_w4a8_convrot_cm2 mul_mat_w4a8_convrot_cm2 ${SD_VK_CM2} -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) diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index 5ac8151854..aad36900be 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -1,6 +1,7 @@ #include #include #include +#include #include #include #include "core/ggml_extensions.h" @@ -12,7 +13,10 @@ #include "mul_mat_i8_tensorwise_cm1.h" #include "mul_mat_i8_tensorwise_cm2.h" #include "mul_mat_w4a4_convrot.h" +#include "mul_mat_w4a4_convrot_cm2.h" +#include "mul_mat_w4a4_convrot_decode_cm2.h" #include "mul_mat_w4a8_convrot.h" +#include "mul_mat_w4a8_convrot_cm2.h" #include "quantize_i8_convrot.h" #include "quantize_i8_convrot_h64.h" #include "quantize_i8_convrot_register.h" @@ -23,11 +27,17 @@ #include "swiglu_quantize_i8_convrot_16384.h" namespace sd_vulkan_i8 { + // Keep these budgets in sync with the cooperative matrix shader tile layouts. + static constexpr uint32_t w4a4_shared = 2 * (64 + 64) * (64 / 4 + 4) * 4; + static constexpr uint32_t w4a4_large_shared = 2 * (64 + 128) * (64 / 4 + 4) * 4; + static constexpr uint32_t w4a8_shared = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 64 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4; + static constexpr uint32_t w4a8_large_shared = (64 * (64 / 2 + 4) + 64 * (64 / 16) + 128 * (64 / 4 + 4) + 256 + 256 + 16 + 1) * 4; 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; + bool w4a4_cm2, w4a8_cm2; uint32_t cm_m, cm_n, cm_k; }; static bool device_info(ggml_backend_dev_t dev, Device& d) { @@ -55,6 +65,10 @@ namespace sd_vulkan_i8 { d.cm2 = true; d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0; } + if (prop.flexible && prop.workgroup_scope && prop.invocations == 256 && d.info.max_group_size[0] >= 256 && 64 % prop.m == 0 && 64 % prop.n == 0) { + d.w4a4_cm2 |= d.compute.shared_memory >= w4a4_shared && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_cm2) > 4; + d.w4a8_cm2 |= d.compute.shared_memory >= w4a8_shared && 32 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_w4a8_convrot_cm2) > 4; + } } return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4; } @@ -698,9 +712,45 @@ namespace sd_vulkan_i8 { struct Parameters { uint32_t n, rows, k, bias; }; - const Parameters pc = {n, rows, k, op.bias ? 1u : 0u}; - auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters); - const uint32_t groups[] = {(n + 63) / 64, (rows + 31) / 32, 1}; + const Parameters pc = {n, rows, k, op.bias ? 1u : 0u}; + auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters); + uint32_t tile_cols = 64, tile_rows = 32; + const bool large_w4a4 = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared; + const bool large_w4a8 = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared; + const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 && + aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4; + const uint32_t constants[] = {64, large_w4a4 || large_w4a8 ? 128u : 64u, 64}; + const char* force_dp4a = std::getenv("GGML_VK_W4_FORCE_DP4A"); + if (!(force_dp4a && std::atoi(force_dp4a)) && aligned(op.input, 16) && + (op.kind == SDW4Kind::W4A4 ? d.w4a4_cm2 : d.w4a8_cm2 && aligned(op.weight, 8))) { + pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot_cm2, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot_cm2, 8, Parameters); + if (decode_w4a4) { + struct DecodeParameters { + uint32_t n, rows, k, bias, row_offset; + }; + DecodeParameters decoded = {n, rows, k, pc.bias, 0}; + uint32_t decode_constants[] = {64, 128, 128, 1}; + pipeline = SD_VK_PIPELINE(mul_mat_w4a4_convrot_decode_cm2, 6, DecodeParameters); + pipeline.constants = decode_constants; + pipeline.constant_count = 4; + uint32_t groups[] = {n / 64, rows / 128, 1}; + auto status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups); + if (status == GGML_STATUS_SUCCESS && rows % 128) { + // Specialize the tail separately so bounds checks do not inhibit full-tile loads. + decoded.row_offset = rows / 128 * 128; + decode_constants[1] = rows % 128 <= 64 ? 64 : 128; + decode_constants[3] = 0; + groups[1] = 1; + status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &decoded, groups); + } + return finish(e, d, output, dst, alias, status); + } + pipeline.constants = constants; + pipeline.constant_count = 3; + tile_cols = constants[0]; + tile_rows = constants[1]; + } + const uint32_t groups[] = {(n + tile_cols - 1) / tile_cols, (rows + tile_rows - 1) / tile_rows, 1}; auto status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups); return finish(e, d, output, dst, alias, status); } diff --git a/src/core/vulkan/shaders/mul_mat_w4a4_convrot_cm2.comp b/src/core/vulkan/shaders/mul_mat_w4a4_convrot_cm2.comp new file mode 100644 index 0000000000..52897e67c8 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a4_convrot_cm2.comp @@ -0,0 +1,88 @@ +#version 460 +#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(local_size_x = 256) in; +layout(constant_id = 0) const uint BM = 64; +layout(constant_id = 1) const uint BN = 64; +layout(constant_id = 2) const uint BK = 64; +const uint STRIDE = BK / 4 + 4; +layout(push_constant) uniform parameter { + uint n; + uint rows; + uint k; + uint has_bias; +} +p; +layout(binding = 0, std430) readonly buffer W { + uint data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + uvec4 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 { + int32_t data_d[]; +}; +shared uint weights[2 * BM * STRIDE]; +shared uint activations[2 * BN * STRIDE]; +void load_tile(uint base, uint buf, uint row_base, uint col_base) { + const uint tid = gl_LocalInvocationIndex; + for (uint i = tid; i < BN * BK / 16; i += 256) { + const uint r = i / (BK / 16), t = i % (BK / 16); + const uvec4 value = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0); + const uint offset = (buf * BN + r) * STRIDE + t * 4; + activations[offset] = value.x; + activations[offset + 1] = value.y; + activations[offset + 2] = value.z; + activations[offset + 3] = value.w; + } + for (uint i = tid; i < BM * BK / 8; i += 256) { + const uint c = i / (BK / 8), t = i % (BK / 8); + const uint v = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 8) + base / 8 + t] : 0u; + const uint lo = (v & 15u) | ((v << 4) & 0xf00u) | ((v << 8) & 0xf0000u) | ((v << 12) & 0xf000000u); + const uint hi = ((v >> 16) & 15u) | ((v >> 12) & 0xf00u) | ((v >> 8) & 0xf0000u) | ((v >> 4) & 0xf000000u); + weights[(buf * BM + c) * STRIDE + 2 * t] = (lo & 0x07070707u) | ((lo & 0x08080808u) * 31u); + weights[(buf * BM + c) * STRIDE + 2 * t + 1] = (hi & 0x07070707u) | ((hi & 0x08080808u) * 31u); + } +} +int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) { + if (row + c >= p.rows || col + r >= p.n) + return 0; + precise float value = float(v) * data_as[row + c] * data_ws[col + r]; + if (p.has_bias != 0) + value += data_b[col + r]; + return floatBitsToInt(value); +} +void main() { + const uint row = gl_WorkGroupID.y * BN, col = gl_WorkGroupID.x * BM; + coopmat sum = coopmat(0); + load_tile(0, 0, row, col); + barrier(); + for (uint base = 0; base < p.k; base += BK) { + const uint buf = (base / BK) & 1u; + if (base + BK < p.k) + load_tile(base + BK, buf ^ 1u, row, col); + coopmat w; + coopmat a; + coopMatLoad(w, weights, buf * BM * STRIDE, STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(a, activations, buf * BN * STRIDE, STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + sum = coopMatMulAdd(w, a, sum); + barrier(); + } + coopMatPerElementNV(sum, sum, dequantize, row, col); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + dl = setTensorLayoutDimensionNV(dl, p.rows, p.n); + tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); + coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose); +} diff --git a/src/core/vulkan/shaders/mul_mat_w4a4_convrot_decode_cm2.comp b/src/core/vulkan/shaders/mul_mat_w4a4_convrot_decode_cm2.comp new file mode 100644 index 0000000000..726c963289 --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a4_convrot_decode_cm2.comp @@ -0,0 +1,92 @@ +#version 460 +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_8bit_storage : require +#extension GL_EXT_buffer_reference : 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(local_size_x = 256) in; +layout(constant_id = 0) const uint BM = 64; +layout(constant_id = 1) const uint BN = 64; +layout(constant_id = 2) const uint BK = 64; +layout(constant_id = 3) const bool FULL = true; +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 { + uvec4 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 { + int32_t data_d[]; +}; +layout(buffer_reference, std430, buffer_reference_align = 16) buffer PackedBlock { + uvec4 value; +}; +int8_t decode_weight(const in PackedBlock block, const in uint block_coords[2], const in uint coord[2]) { + return int8_t(bitfieldExtract(int(block.value[coord[1] / 8]), int((coord[1] % 8) * 4), 4)); +} +int32_t dequantize(const in uint r, const in uint c, const in int32_t v, const in uint row, const in uint col) { + if (row + c >= p.rows || col + r >= p.n) + return 0; + precise float value = float(v) * data_as[row + c] * data_ws[col + r]; + if (p.has_bias != 0) + value += data_b[col + r]; + return floatBitsToInt(value); +} +// The host selects complete weight tiles and a 256-aligned input width. +void main() { + const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y; + const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x; + const uint panel = min(32u, cols), first_col = (linear / (panel * rows)) * panel; + const uint group_cols = min(panel, cols - first_col), within = linear - first_col * rows; + const uint row = p.row_offset + (within / group_cols) * BN, col = (first_col + within % group_cols) * BM; + coopmat sum = coopmat(0); + tensorLayoutNV<2> wl = createTensorLayoutNV(2); + wl = setTensorLayoutBlockSizeNV(wl, 1, 32); + wl = setTensorLayoutDimensionNV(wl, p.n, p.k & ~255u); + tensorLayoutNV<2> al = createTensorLayoutNV(2); + al = setTensorLayoutDimensionNV(al, p.rows, p.k & ~255u); + tensorViewNV<2, false, 1, 0> at = createTensorViewNV(2, false, 1, 0); + if (FULL) { + for (uint base = 0; base < p.k; base += BK) { + coopmat w; + coopmat a; + coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(wl, col, BM, base, BK), decode_weight); + coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(al, row, BN, base, BK), at); + sum = coopMatMulAdd(w, a, sum); + } + } else { + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tail = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + tail = setTensorLayoutDimensionNV(tail, p.rows, p.k & ~255u); + for (uint base = 0; base < p.k; base += BK) { + coopmat w; + coopmat a; + coopMatLoadTensorNV(w, data_w, 0, sliceTensorLayoutNV(wl, col, BM, base, BK), decode_weight); + coopMatLoadTensorNV(a, data_a, 0, sliceTensorLayoutNV(tail, row, BN, base, BK), at); + sum = coopMatMulAdd(w, a, sum); + } + } + coopMatPerElementNV(sum, sum, dequantize, row, col); + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> dl = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + dl = setTensorLayoutDimensionNV(dl, p.rows, p.n); + tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); + coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(dl, row, BN, col, BM), transpose); +} diff --git a/src/core/vulkan/shaders/mul_mat_w4a8_convrot_cm2.comp b/src/core/vulkan/shaders/mul_mat_w4a8_convrot_cm2.comp new file mode 100644 index 0000000000..c97f134e4e --- /dev/null +++ b/src/core/vulkan/shaders/mul_mat_w4a8_convrot_cm2.comp @@ -0,0 +1,164 @@ +#version 460 + +#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 +#extension GL_EXT_control_flow_attributes : require + +layout(local_size_x = 256) in; +layout(constant_id = 0) const uint BM = 64; +layout(constant_id = 1) const uint BN = 64; +layout(constant_id = 2) const uint BK = 64; +const uint STRIDE = BK / 2 + 4; + +layout(push_constant) uniform parameter { + uint n; + uint rows; + uint k; + uint has_bias; +} +p; + +layout(binding = 0, std430) readonly buffer W { + uvec2 data_w[]; +}; +layout(binding = 1, std430) readonly buffer A { + uvec4 data_a[]; +}; +layout(binding = 2, std430) readonly buffer CB { + float data_cb[]; +}; +layout(binding = 3, std430) readonly buffer SC { + float data_sc[]; +}; +layout(binding = 4, std430) readonly buffer SR { + uint data_sr[]; +}; +layout(binding = 5, std430) readonly buffer AS { + float data_as[]; +}; +layout(binding = 6, std430) readonly buffer B { + float data_b[]; +}; +layout(binding = 7, std430) writeonly buffer D { + int32_t data_d[]; +}; +shared uint lut[256]; +shared int cb[16]; +shared float cb_scale; +shared float relative[256]; +shared float scales[BK / 16 * BM]; + +float fp8_e4m3(uint v) { + const uint exponent = (v >> 3) & 15u, mantissa = v & 7u; + const float value = exponent == 15u && mantissa == 7u ? uintBitsToFloat(0x7fc00000u) : exponent == 0u ? ldexp(float(mantissa), -9) + : ldexp(1.0 + float(mantissa) / 8.0, int(exponent) - 7); + return (v & 128u) != 0u ? -value : value; +} + +int32_t accumulate_group(const in uint row, const in uint col, const in int32_t value, const in int32_t previous, const in uint group) { + return floatBitsToInt(fma(float(value), scales[group * BM + row], intBitsToFloat(previous))); +} + +shared uint weights[BM * STRIDE]; +const uint A_STRIDE = BK / 4 + 4; +shared uint activations[BN * A_STRIDE]; + +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, c = col_base + row; + if (r >= p.rows || c >= p.n) { + return 0; + } + precise float value = intBitsToFloat(acc) * data_as[r]; + if (p.has_bias != 0) { + value += data_b[c]; + } + return floatBitsToInt(value); +} + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint cols = gl_NumWorkGroups.x, rows = gl_NumWorkGroups.y; + const uint linear = gl_WorkGroupID.y * cols + gl_WorkGroupID.x; + // Reuse weight panels across activation tiles when the output expands. + const uint panel = p.n > p.k ? 16u : cols; + const uint first_col = (linear / (panel * rows)) * panel; + const uint group_cols = min(panel, cols - first_col); + const uint within = linear - first_col * rows; + const uint col_base = (first_col + within % group_cols) * BM, row_base = (within / group_cols) * BN; + tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> d_layout = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); + d_layout = setTensorLayoutDimensionNV(d_layout, p.rows, p.n); + tensorViewNV<2, false, 1, 0> transpose = createTensorViewNV(2, false, 1, 0); + coopmat sum = coopmat(0); + if (tid == 0) { + float amax = 0.0; + for (uint i = 0; i < 16; ++i) { + amax = max(amax, abs(data_cb[i])); + } + cb_scale = amax / 127.0; + for (uint i = 0; i < 16; ++i) { + cb[i] = cb_scale == 0.0 ? 0 : clamp(int(roundEven(data_cb[i] / cb_scale)), -127, 127); + } + } + relative[tid] = fp8_e4m3(tid); + barrier(); + for (uint i = tid; i < 256; i += 256) { + lut[i] = uint(cb[i & 15u] & 255) | (uint(cb[i >> 4] & 255) << 8); + } + barrier(); + for (uint base = 0; base < p.k; base += BK) { + for (uint i = tid; i < BN * BK / 16; i += 256) { + const uint r = i / (BK / 16), t = i % (BK / 16); + const uvec4 value = row_base + r < p.rows ? data_a[(row_base + r) * (p.k / 16) + base / 16 + t] : uvec4(0); + const uint offset = r * A_STRIDE + t * 4; + activations[offset] = value.x; + activations[offset + 1] = value.y; + activations[offset + 2] = value.z; + activations[offset + 3] = value.w; + } + for (uint i = tid; i < BM * BK / 16; i += 256) { + const uint c = i / (BK / 16), part = i % (BK / 16); + const uvec2 packed = col_base + c < p.n ? data_w[(col_base + c) * (p.k / 16) + base / 16 + part] : uvec2(0); + // Complementary zero halves let adjacent groups reuse one activation fragment. + const uint odd = part % 2; + const uint offset = c * STRIDE + part * 8 + odd * 4; + const uint zero_offset = c * STRIDE + part * 8 + (1 - odd) * 4; + for (uint j = 0; j < 2; ++j) { + const uint v = packed[j]; + weights[offset + 2 * j] = lut[v & 255u] | (lut[(v >> 8) & 255u] << 16); + weights[offset + 2 * j + 1] = lut[(v >> 16) & 255u] | (lut[v >> 24] << 16); + weights[zero_offset + 2 * j] = 0u; + weights[zero_offset + 2 * j + 1] = 0u; + } + } + for (uint i = tid; i < BM * BK / 16; i += 256) { + const uint c = i / (BK / 16), g = i % (BK / 16); + float scale = 0.0; + if (col_base + c < p.n) { + const uint packed = data_sr[(col_base + c) * (p.k / 64) + base / 64 + g / 4]; + scale = cb_scale * (data_sc[col_base + c] * relative[(packed >> (8 * (g % 4))) & 255u]); + } + scales[g * BM + c] = scale; + } + barrier(); + // The relative scale changes every 16 input values, so each group + // needs its own integer MMA and floating-point accumulation. Pad the + // weights with zeros to satisfy the device's K=32 matrix granularity. + [[unroll]] for (uint g = 0; g < BK / 16; g += 2) { + coopmat a; + coopMatLoad(a, activations, g * 4, A_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + [[unroll]] for (uint j = 0; j < 2; ++j) { + coopmat w; + coopmat part = coopmat(0); + coopMatLoad(w, weights, (g + j) * 8, STRIDE, gl_CooperativeMatrixLayoutRowMajor); + part = coopMatMulAdd(w, a, part); + coopMatPerElementNV(sum, part, accumulate_group, sum, g + j); + } + } + barrier(); + } + coopMatPerElementNV(sum, sum, dequantize, row_base, col_base); + coopMatStoreTensorNV(sum, data_d, 0, sliceTensorLayoutNV(d_layout, row_base, BN, col_base, BM), transpose); +} From fdfc913c515ad749cedd84a118e7c23a05ac1c65 Mon Sep 17 00:00:00 2001 From: Daniele <57776841+daniandtheweb@users.noreply.github.com> Date: Sat, 7 Nov 2026 15:33:41 +0000 Subject: [PATCH 23/57] feat: ControlNet Union support for SDXL (#2284) --- examples/common/common.cpp | 19 +++ examples/common/common.h | 2 + examples/server/api.md | 3 + examples/server/routes_sdcpp.cpp | 1 + include/stable-diffusion.h | 18 +++ src/generation_info.cpp | 7 +- src/model/common/block.hpp | 27 +++++ src/model/diffusion/control.hpp | 195 +++++++++++++++++++++++++++++- src/model/diffusion/unet.hpp | 6 +- src/name_conversion.cpp | 9 ++ src/pipeline/diffusion_engine.cpp | 4 + src/pipeline/diffusion_engine.h | 4 + src/pipeline/image.cpp | 3 + src/pipeline/request.cpp | 1 + src/pipeline/request.h | 1 + src/stable-diffusion.cpp | 53 ++++++++ 16 files changed, 343 insertions(+), 10 deletions(-) diff --git a/examples/common/common.cpp b/examples/common/common.cpp index d5496f66fe..02c30880af 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1187,6 +1187,13 @@ ArgOptions SDGenerationParams::get_options() { "under --hires-upscalers-dir (default: Latent)", 0, &hires_upscaler}, + {"", + "--control-type", + "kind of the control image, one of [openpose, depth, scribble, canny, normal, segment, tile, repaint]. " + "Required by an SDXL Control Net Union model (tile and repaint only exist in its ProMax variant), ignored by any other Control Net. " + "Also accepted: hed, softedge, pidi, ted (scribble); lineart, anime_lineart, mlsd (canny); seg (segment); inpaint, outpaint (repaint)", + 0, + &control_type}, {"", "--extra-sample-args", "extra sampler/scheduler/guidance args, key=value list. CFG supports guidance_schedule; APG supports apg_eta, apg_momentum, apg_norm_threshold, apg_norm_threshold_smoothing; SLG supports slg_uncond; lcm supports noise_clip_std, noise_scale_start, noise_scale_end; flux supports base_shift, max_shift; ltx2 supports max_shift, base_shift, stretch, terminal; euler_ge supports gamma; beta scheduler supports alpha, beta; logit_normal supports mu, std, logsnr_min, logsnr_max, resolution_aware; llada_image supports uniform; lms supports lms_max_order, lms_shift, lms_divisions; noise-injecting samplers support noise_sampler with value iid (default except for dpm++2m_sde_bt) or brownian_tree; brownian_tree_rng supports cpu (default), cuda, std_default or sampler_rng", @@ -2249,6 +2256,7 @@ bool SDGenerationParams::from_json_str( load_if_exists("strength", strength); load_if_exists("control_strength", control_strength); + load_if_exists("control_type", control_type); load_if_exists("ip_adapter_strength", ip_adapter_strength); load_if_exists("moe_boundary", moe_boundary); load_if_exists("vace_strength", vace_strength); @@ -2638,6 +2646,15 @@ bool SDGenerationParams::resolve(const std::string& lora_model_dir, const std::s sample_params.sample_steps = std::clamp(sample_params.sample_steps, 1, 100); } + resolved_control_type = SD_CONTROL_TYPE_NONE; + if (!control_type.empty()) { + resolved_control_type = str_to_sd_control_type(control_type.c_str()); + if (resolved_control_type == SD_CONTROL_TYPE_COUNT) { + LOG_ERROR("invalid control type '%s' (openpose, depth, scribble, canny, normal, segment, tile, repaint)", control_type.c_str()); + return false; + } + } + hires_upscaler_model_path.clear(); if (hires_enabled) { if (hires_upscaler.empty()) { @@ -2877,6 +2894,7 @@ sd_img_gen_params_t SDGenerationParams::to_sd_img_gen_params_t() { params.qwen_image_layers = qwen_image_layers; params.control_image = control_image.get(); params.control_strength = control_strength; + params.control_type = resolved_control_type; params.ip_adapter_image = ip_adapter_image.get(); params.ip_adapter_strength = ip_adapter_strength; params.pm_params = pm_params; @@ -3082,6 +3100,7 @@ std::string SDGenerationParams::to_string() const { << " vace_strength: " << vace_strength << ",\n" << " strength: " << strength << ",\n" << " control_strength: " << control_strength << ",\n" + << " control_type: \"" << control_type << "\",\n" << " seed: " << seed << ",\n" << " upscale_repeats: " << upscale_repeats << ",\n" << " upscale_tile_size: " << upscale_tile_size << ",\n" diff --git a/examples/common/common.h b/examples/common/common.h index 39d6011afc..4ecbef4fa3 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -227,6 +227,7 @@ struct SDGenerationParams { std::string end_image_path; std::string mask_image_path; std::string control_image_path; + std::string control_type; std::string ip_adapter_image_path; std::vector ref_image_paths; std::vector ref_video_paths; @@ -291,6 +292,7 @@ struct SDGenerationParams { std::string prompt_with_lora; // for metadata record only std::vector lora_vec; sd_hires_upscaler_t resolved_hires_upscaler; + sd_control_type_t resolved_control_type = SD_CONTROL_TYPE_NONE; // Owned execution payload. SDImageOwner init_image; diff --git a/examples/server/api.md b/examples/server/api.md index 97cbd66cdf..9c8d8c8ce9 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -936,6 +936,7 @@ LTX and Wan preserve causal state between temporal tiles. Hunyuan Video and TAEH | `ref_image_args` | `string` | | `increase_ref_index` | `boolean` | | `control_strength` | `number` | +| `control_type` | `string` | | `ip_adapter_strength` | `number` | | `hires` | `object` | | `hires.enabled` | `boolean` | @@ -1116,6 +1117,7 @@ Example: "ref_image_args": "", "increase_ref_index": false, "control_strength": 0.9, + "control_type": "", "ip_adapter_strength": 1.0, "embed_image_metadata": true, @@ -1241,6 +1243,7 @@ Top-level scalar fields: | `image_preprocess` | `string \| array` | | `increase_ref_index` | `boolean` | | `control_strength` | `number` | +| `control_type` | `string` | | `ip_adapter_strength` | `number` | | `embed_image_metadata` | `boolean` | | `preview` | `string` | diff --git a/examples/server/routes_sdcpp.cpp b/examples/server/routes_sdcpp.cpp index 0dda9a95be..1d11b89adc 100644 --- a/examples/server/routes_sdcpp.cpp +++ b/examples/server/routes_sdcpp.cpp @@ -154,6 +154,7 @@ static json make_img_gen_defaults_json(const SDGenerationParams& defaults, const {"image_preprocess", defaults.image_preprocess}, {"increase_ref_index", defaults.increase_ref_index}, {"control_strength", defaults.control_strength}, + {"control_type", defaults.control_type}, {"ip_adapter_strength", defaults.ip_adapter_strength}, {"sample_params", make_sample_params_json(defaults.sample_params, defaults.skip_layers, defaults.custom_sigmas)}, {"hires", make_hires_json(defaults)}, diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index f4ea0eaf13..1c61e4b150 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -400,6 +400,19 @@ enum sd_hires_upscaler_t { SD_HIRES_UPSCALER_COUNT, }; +enum sd_control_type_t { + SD_CONTROL_TYPE_NONE, + SD_CONTROL_TYPE_OPENPOSE, + SD_CONTROL_TYPE_DEPTH, + SD_CONTROL_TYPE_SCRIBBLE, // thick lines: scribble, hed, softedge, ted + SD_CONTROL_TYPE_CANNY, // thin lines: canny, lineart, anime lineart, mlsd + SD_CONTROL_TYPE_NORMAL, + SD_CONTROL_TYPE_SEGMENT, + SD_CONTROL_TYPE_TILE, + SD_CONTROL_TYPE_REPAINT, // inpaint, outpaint + SD_CONTROL_TYPE_COUNT, +}; + typedef struct { bool enabled; enum sd_hires_upscaler_t upscaler; @@ -444,6 +457,7 @@ typedef struct { bool circular_x; bool circular_y; sd_image_preprocess_params_t image_preprocess; + enum sd_control_type_t control_type; } sd_img_gen_params_t; typedef struct { @@ -531,6 +545,10 @@ SD_API const char* sd_lora_apply_mode_name(enum lora_apply_mode_t mode); SD_API enum lora_apply_mode_t str_to_lora_apply_mode(const char* str); SD_API const char* sd_hires_upscaler_name(enum sd_hires_upscaler_t upscaler); SD_API enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str); +SD_API const char* sd_control_type_name(enum sd_control_type_t control_type); +// Also accepts the names of the kinds a type stands for (hed, lineart, mlsd, ...); +// returns SD_CONTROL_TYPE_COUNT when the name is unknown. +SD_API enum sd_control_type_t str_to_sd_control_type(const char* str); SD_API void sd_cache_params_init(sd_cache_params_t* cache_params); SD_API void sd_hires_params_init(sd_hires_params_t* hires_params); diff --git a/src/generation_info.cpp b/src/generation_info.cpp index 2f49054ea1..cf40bc2bad 100644 --- a/src/generation_info.cpp +++ b/src/generation_info.cpp @@ -127,8 +127,11 @@ namespace sd { auto info = common_info(p, "image"); if (p.init_image.data) info["strength"] = p.strength; - if (p.control_image.data) + if (p.control_image.data) { info["control_strength"] = p.control_strength; + if (p.control_type != SD_CONTROL_TYPE_NONE) + info["control_type"] = sd_control_type_name(p.control_type); + } 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) { @@ -188,7 +191,7 @@ namespace sd { 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"}) + for (const auto* key : {"control_strength", "control_type", "pm_params", "pulid_params", "hires", "image_preprocess"}) inpaint.erase(key); sd_img_gen_params_t defaults; sd_img_gen_params_init(&defaults); diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp index 66180fa8d1..5be333d6fa 100644 --- a/src/model/common/block.hpp +++ b/src/model/common/block.hpp @@ -576,6 +576,33 @@ class SpatialTransformer : public GGMLBlock { } }; +class ResidualAttentionBlock : public GGMLBlock { +public: + ResidualAttentionBlock(int64_t channels, int64_t n_head) { + blocks["ln_1"] = std::shared_ptr(new LayerNorm(channels)); + blocks["attn"] = std::shared_ptr(new MultiheadAttention(channels, n_head, true, true, true)); + blocks["ln_2"] = std::shared_ptr(new LayerNorm(channels)); + blocks["mlp.c_fc"] = std::shared_ptr(new Linear(channels, channels * 4)); + blocks["mlp.c_proj"] = std::shared_ptr(new Linear(channels * 4, channels)); + } + + // x: [N, n_token, channels] + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto ln_1 = std::dynamic_pointer_cast(blocks["ln_1"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto ln_2 = std::dynamic_pointer_cast(blocks["ln_2"]); + auto mlp_c_fc = std::dynamic_pointer_cast(blocks["mlp.c_fc"]); + auto mlp_c_proj = std::dynamic_pointer_cast(blocks["mlp.c_proj"]); + + x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, ln_1->forward(ctx, x))); + + auto h = mlp_c_fc->forward(ctx, ln_2->forward(ctx, x)); + h = ggml_ext_gelu_quick(ctx->ggml_ctx, h); + h = mlp_c_proj->forward(ctx, h); + return ggml_add(ctx->ggml_ctx, x, h); + } +}; + class AlphaBlender : public GGMLBlock { protected: void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, std::string prefix = "") override { diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp index 3b08ed1500..7848892d0f 100644 --- a/src/model/diffusion/control.hpp +++ b/src/model/diffusion/control.hpp @@ -2,10 +2,47 @@ #define __SD_MODEL_DIFFUSION_CONTROL_HPP__ #include "model/common/block.hpp" +#include "stable-diffusion.h" // Match main UNet's MAX_GRAPH_SIZE so SDXL ControlNet (transformer_depth={1,2,10}) fits. #define CONTROL_NET_GRAPH_SIZE MAX_GRAPH_SIZE +/* + =============================== ControlNet Union ================================ + Reference: https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/controlnets/controlnet_union.py + +*/ +struct ControlNetUnionConfig { + int num_control_types = 0; // 0: not a union model + int num_fusion_layers = 0; + + static constexpr int CONTROL_TYPE_EMBED_DIM = 256; + static constexpr int FUSION_HEADS = 8; + + static ControlNetUnionConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix, + SDVersion version = VERSION_SD1) { + ControlNetUnionConfig config; + if (!sd_version_is_sdxl(version)) { + return config; + } + const std::string p = prefix.empty() ? "" : prefix + "."; + auto embedding = tensor_storage_map.find(p + "control_add_embedding.linear_1.weight"); + if (embedding == tensor_storage_map.end()) { + return config; + } + // control_add_embedding takes CONTROL_TYPE_EMBED_DIM inputs per control type + config.num_control_types = static_cast(embedding->second.logical_ne0() / CONTROL_TYPE_EMBED_DIM); + while (tensor_storage_map.find(p + "transformer_layes." + std::to_string(config.num_fusion_layers) + ".ln_1.weight") != tensor_storage_map.end()) { + config.num_fusion_layers++; + } + LOG_VERBOSE("controlnet union: num_control_types = %d, num_fusion_layers = %d", + config.num_control_types, + config.num_fusion_layers); + return config; + } +}; + /* =================================== ControlNet =================================== Reference: https://github.com/comfyanonymous/ComfyUI/blob/master/comfy/cldm/cldm.py @@ -13,6 +50,8 @@ */ class ControlNetBlock : public GGMLBlock { protected: + ControlNetUnionConfig config; + SDVersion version = VERSION_SD1; // network hparams int in_channels = 4; @@ -28,12 +67,19 @@ class ControlNetBlock : public GGMLBlock { int context_dim = 768; // 1024 for VERSION_SD2, 2048 for VERSION_SDXL bool use_linear_projection = false; + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + if (is_union()) { + // [num_control_types, model_channels] in the checkpoint + params["task_embedding"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, model_channels, config.num_control_types); + } + } + public: int model_channels = 320; int adm_in_channels = 2816; // only for VERSION_SDXL - ControlNetBlock(SDVersion version = VERSION_SD1) - : version(version) { + ControlNetBlock(SDVersion version = VERSION_SD1, ControlNetUnionConfig config = {}) + : config(config), version(version) { if (sd_version_is_sd2(version)) { context_dim = 1024; num_head_channels = 64; @@ -162,6 +208,76 @@ class ControlNetBlock : public GGMLBlock { // middle_block_out blocks["middle_block_out.0"] = std::shared_ptr(make_zero_conv(ch)); + + if (is_union()) { + blocks["control_add_embedding.linear_1"] = std::shared_ptr(new Linear(ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, time_embed_dim)); + // nn.SiLU() + blocks["control_add_embedding.linear_2"] = std::shared_ptr(new Linear(time_embed_dim, time_embed_dim)); + for (int i = 0; i < config.num_fusion_layers; i++) { + blocks["transformer_layes." + std::to_string(i)] = std::shared_ptr(new ResidualAttentionBlock(model_channels, ControlNetUnionConfig::FUSION_HEADS)); + } + blocks["spatial_ch_projs"] = std::shared_ptr(new Linear(model_channels, model_channels)); + } + } + + bool is_union() const { + return config.num_control_types > 0; + } + + // control_types: [N*num_control_types], 1 for the types in use and 0 for the others + ggml_tensor* control_type_embedding(GGMLRunnerContext* ctx, ggml_tensor* control_types, int64_t batch) { + auto linear_1 = std::dynamic_pointer_cast(blocks["control_add_embedding.linear_1"]); + auto linear_2 = std::dynamic_pointer_cast(blocks["control_add_embedding.linear_2"]); + + // [N*num_control_types, CONTROL_TYPE_EMBED_DIM] + auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM); + e = ggml_reshape_2d(ctx->ggml_ctx, e, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, batch); + e = linear_1->forward(ctx, e); + e = ggml_silu_inplace(ctx->ggml_ctx, e); + return linear_2->forward(ctx, e); // [N, time_embed_dim] + } + + // [N, C, H, W] => [N, C] + static ggml_tensor* spatial_mean(ggml_context* ctx, ggml_tensor* x) { + auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); + return ggml_reshape_2d(ctx, ggml_mean(ctx, flat), x->ne[2], x->ne[3]); + } + + // sample: [N, C, H, W], hint: [N, C, H, W] or [1, C, H, W] + ggml_tensor* fuse_condition(GGMLRunnerContext* ctx, + ggml_tensor* sample, + ggml_tensor* hint, + int control_type_index) { + auto c = ctx->ggml_ctx; + int64_t channels = sample->ne[2]; + int64_t batch = sample->ne[3]; + + auto task_embedding = params["task_embedding"]; // [num_control_types, C] + auto task = ggml_view_2d(c, task_embedding, channels, 1, task_embedding->nb[1], control_type_index * task_embedding->nb[1]); + + auto hint_feat = ggml_add(c, spatial_mean(c, hint), task); // [N, C] + if (hint_feat->ne[1] != batch) { + hint_feat = ggml_repeat(c, hint_feat, ggml_new_tensor_2d(c, GGML_TYPE_F32, channels, batch)); + } + auto sample_feat = spatial_mean(c, sample); // [N, C] + + // as in the reference (nn.MultiheadAttention without batch_first) the attention is over + // the batch, and the 2 features (hint, sample) are its batch: [2, N, C] + auto x = ggml_concat(c, + ggml_reshape_3d(c, hint_feat, channels, batch, 1), + ggml_reshape_3d(c, sample_feat, channels, batch, 1), + 2); + for (int i = 0; i < config.num_fusion_layers; i++) { + auto layer = std::dynamic_pointer_cast(blocks["transformer_layes." + std::to_string(i)]); + x = layer->forward(ctx, x); + } + + auto projection = std::dynamic_pointer_cast(blocks["spatial_ch_projs"]); + auto alpha = ggml_reshape_2d(c, ggml_ext_slice(c, x, 2, 0, 1), channels, batch); + alpha = projection->forward(ctx, alpha); + alpha = ggml_reshape_4d(c, alpha, 1, 1, channels, batch); + + return ggml_add(c, ggml_add(c, sample, hint), alpha); } ggml_tensor* resblock_forward(std::string name, @@ -204,7 +320,9 @@ class ControlNetBlock : public GGMLBlock { ggml_tensor* guided_hint, ggml_tensor* timesteps, ggml_tensor* context, - ggml_tensor* y = nullptr) { + ggml_tensor* y = nullptr, + ggml_tensor* control_types = nullptr, + int control_type_index = 0) { // x: [N, in_channels, h, w] or [N, in_channels/2, h, w] // timesteps: [N,] // context: [N, max_position, hidden_size] or [1, max_position, hidden_size]. for example, [N, 77, 768] @@ -246,6 +364,11 @@ class ControlNetBlock : public GGMLBlock { emb = ggml_add(ctx->ggml_ctx, emb, label_emb); // [N, time_embed_dim] } + if (is_union()) { + GGML_ASSERT(control_types != nullptr); + emb = ggml_add(ctx->ggml_ctx, emb, control_type_embedding(ctx, control_types, x->ne[3])); + } + std::vector outs; if (guided_hint == nullptr) { @@ -257,7 +380,11 @@ class ControlNetBlock : public GGMLBlock { // input block 0 auto h = input_blocks_0_0->forward(ctx, x); - h = ggml_add(ctx->ggml_ctx, h, guided_hint); + if (is_union()) { + h = fuse_condition(ctx, h, guided_hint, control_type_index); + } else { + h = ggml_add(ctx->ggml_ctx, h, guided_hint); + } outs.push_back(zero_convs_0->forward(ctx, h)); // input block 1-11 @@ -308,9 +435,15 @@ class ControlNetBlock : public GGMLBlock { struct ControlNet : public GGMLRunner { SDVersion version = VERSION_SD1; + ControlNetUnionConfig config; ControlNetBlock control_net; std::string weight_prefix; + // ControlNet Union: index of the control type of the control image (-1: none yet) + // and the 0/1 vector of the types in use that the graph reads + int control_type_index = -1; + sd::Tensor control_type_tensor; + std::vector control_outputs_ggml; ggml_tensor* guided_hint_output_ggml = nullptr; std::vector> controls; @@ -320,12 +453,46 @@ struct ControlNet : public GGMLRunner { return "controlnet.guided_hint"; } + bool is_union() const { + return config.num_control_types > 0; + } + + // set the kind of the control image of the next generations, false (after saying why) when it + // can't be used: a union model needs one of the types it was trained on + bool set_control_type(sd_control_type_t type) { + if (!is_union()) { + if (type != SD_CONTROL_TYPE_NONE) { + LOG_WARN("control type '%s' ignored: the ControlNet is not a union model", sd_control_type_name(type)); + } + control_type_index = -1; + return true; + } + if (type == SD_CONTROL_TYPE_NONE) { + LOG_ERROR("this ControlNet is a union model and needs a control type (openpose, depth, scribble, canny, normal, segment%s)", + config.num_control_types > 6 ? ", tile, repaint" : ""); + return false; + } + // the types of a union model, in training order: the enum without NONE + const int index = static_cast(type) - 1; + if (index < 0 || index >= config.num_control_types) { + LOG_ERROR("control type '%s' is not supported by this ControlNet Union (it has %d control types)", + sd_control_type_name(type), config.num_control_types); + return false; + } + control_type_index = index; + return true; + } + ControlNet(ggml_backend_t backend, const String2TensorStorage& tensor_storage_map = {}, SDVersion version = VERSION_SD1, const std::string& prefix = "", std::shared_ptr weight_manager = nullptr) - : GGMLRunner(backend, weight_manager), version(version), control_net(version), weight_prefix(prefix) { + : GGMLRunner(backend, weight_manager), + version(version), + config(ControlNetUnionConfig::detect_from_weights(tensor_storage_map, prefix, version)), + control_net(version, config), + weight_prefix(prefix) { control_net.init(params_ctx, tensor_storage_map, prefix); } @@ -376,6 +543,16 @@ struct ControlNet : public GGMLRunner { hint = make_input(hint_tensor); } + ggml_tensor* control_types = nullptr; + if (is_union()) { + const int64_t batch = x_tensor.dim() >= 4 ? x_tensor.shape()[3] : 1; + control_type_tensor = sd::zeros({config.num_control_types * batch}); + for (int64_t i = 0; i < batch; i++) { + control_type_tensor[i * config.num_control_types + control_type_index] = 1.f; + } + control_types = make_input(control_type_tensor); + } + auto runner_ctx = get_context(gf); auto outs = control_net.forward(&runner_ctx, @@ -384,7 +561,9 @@ struct ControlNet : public GGMLRunner { guided_hint_input, timesteps, context, - y); + y, + control_types, + control_type_index); if (guided_hint_input == nullptr && !outs.empty()) { guided_hint_output_ggml = outs[0]; @@ -413,6 +592,10 @@ struct ControlNet : public GGMLRunner { // timesteps: [N, ] // context: [N, max_position, hidden_size]([N, 77, 768]) or [1, max_position, hidden_size] // y: [N, adm_in_channels] or [1, adm_in_channels] + if (is_union() && control_type_index < 0) { + LOG_ERROR("ControlNet Union: no control type set"); + return std::nullopt; + } auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x, hint, timesteps, context, y); }; diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp index 675d94296c..837163fd90 100644 --- a/src/model/diffusion/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -684,7 +684,8 @@ class UnetModelBlock : public GGMLBlock { } sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h"); if (controls.size() > 0) { - auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, true); + // not in place, it segfaults the Vulkan backend with SDXL ControlNets (seen on RADV) + auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, false); h = ggml_add(ctx->ggml_ctx, h, cs); // middle control } int control_offset = static_cast(controls.size() - 2); @@ -697,7 +698,8 @@ class UnetModelBlock : public GGMLBlock { hs.pop_back(); if (controls.size() > 0) { - auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, true); + // not in place, see the middle control + auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, false); h_skip = ggml_add(ctx->ggml_ctx, h_skip, cs); // control net condition control_offset--; } diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 2e147348ff..1f19b11879 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1840,6 +1840,15 @@ std::string convert_tensor_name(std::string name, SDVersion version) { if (sd_version_is_sdxl(version) && is_diffusers_controlnet_name(name)) { name = convert_diffusers_controlnet_to_original_sdxl(name); } + // controlnet union fusion: nn.MultiheadAttention has in_proj_weight/in_proj_bias, + // MultiheadAttention here has in_proj.weight/in_proj.bias + if (starts_with(name, "transformer_layes.")) { + const std::string from = "attn.in_proj_"; + size_t pos = name.find(from); + if (pos != std::string::npos) { + name.replace(pos, from.size(), "attn.in_proj."); + } + } } static const std::vector> generic_name_map = { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 4875858b82..a8eed8e771 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2270,6 +2270,10 @@ static sd::Tensor batch_two_condition_tensors(const sd::Tensor& a, return sd::ops::concat(a, b, static_cast(batch_dim)); } +bool StableDiffusionGGML::set_control_type(sd_control_type_t type) { + return control_net == nullptr || control_net->set_control_type(type); +} + void StableDiffusionGGML::compute_sample_controls(const sd::Tensor& control_image, const sd::Tensor& noised_input, const sd::Tensor& timesteps_tensor, diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index e067781b45..97e3acac9b 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -440,6 +440,10 @@ class StableDiffusionGGML { bool terminal_sigma_is_zero, int64_t* last_progress_us); + // set the kind of the control image of the next generations (ControlNet Union), + // false (after saying why) when the loaded ControlNet can't use it + bool set_control_type(sd_control_type_t type); + void compute_sample_controls(const sd::Tensor& control_image, const sd::Tensor& noised_input, const sd::Tensor& timesteps_tensor, diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 0763e38ca0..03a6707d40 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -253,6 +253,9 @@ namespace sd::pipeline { if (sd_img_gen_params->control_image.data != nullptr) { control_image_tensor = sd_image_to_tensor(sd_img_gen_params->control_image, request->width, request->height); + if (!sd->set_control_type(request->control_type)) { + return std::nullopt; + } } if (init_image_tensor.empty() || mask_image_tensor.empty()) { diff --git a/src/pipeline/request.cpp b/src/pipeline/request.cpp index d333fbb8be..739b8e43a1 100644 --- a/src/pipeline/request.cpp +++ b/src/pipeline/request.cpp @@ -135,6 +135,7 @@ namespace sd::pipeline { shifted_timestep = sd_img_gen_params->sample_params.shifted_timestep; strength = sd_img_gen_params->strength; control_strength = sd_img_gen_params->control_strength; + control_type = sd_img_gen_params->control_type; eta = sd_img_gen_params->sample_params.eta; has_ref_images = sd_img_gen_params->ref_images_count > 0; guidance = sd_img_gen_params->sample_params.guidance; diff --git a/src/pipeline/request.h b/src/pipeline/request.h index 1cf1cee6a3..b1ab6404cb 100644 --- a/src/pipeline/request.h +++ b/src/pipeline/request.h @@ -47,6 +47,7 @@ namespace sd::pipeline { int shifted_timestep = 0; float strength = 1.f; float control_strength = 0.f; + sd_control_type_t control_type = SD_CONTROL_TYPE_NONE; float eta = 0.f; sd_sample_params_t sample_params = {}; sd_sample_params_t high_noise_sample_params = {}; diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index bca4b79bac..436328d77e 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -267,6 +267,57 @@ enum sd_hires_upscaler_t str_to_sd_hires_upscaler(const char* str) { return SD_HIRES_UPSCALER_COUNT; } +const char* control_type_to_str[] = { + "none", + "openpose", + "depth", + "scribble", + "canny", + "normal", + "segment", + "tile", + "repaint", +}; + +static_assert(SD_CONTROL_TYPE_COUNT == sizeof(control_type_to_str) / sizeof(control_type_to_str[0]), + "\nnumber of elements in control_type_to_str[] != SD_CONTROL_TYPE_COUNT"); + +const char* sd_control_type_name(enum sd_control_type_t control_type) { + if (control_type >= SD_CONTROL_TYPE_NONE && control_type < SD_CONTROL_TYPE_COUNT) { + return control_type_to_str[control_type]; + } + return NONE_STR; +} + +enum sd_control_type_t str_to_sd_control_type(const char* str) { + for (int i = 0; i < SD_CONTROL_TYPE_COUNT; i++) { + if (!strcmp(str, control_type_to_str[i])) { + return (enum sd_control_type_t)i; + } + } + static const struct { + const char* alias; + enum sd_control_type_t type; + } aliases[] = { + {"hed", SD_CONTROL_TYPE_SCRIBBLE}, + {"softedge", SD_CONTROL_TYPE_SCRIBBLE}, + {"pidi", SD_CONTROL_TYPE_SCRIBBLE}, + {"ted", SD_CONTROL_TYPE_SCRIBBLE}, + {"lineart", SD_CONTROL_TYPE_CANNY}, + {"anime_lineart", SD_CONTROL_TYPE_CANNY}, + {"mlsd", SD_CONTROL_TYPE_CANNY}, + {"seg", SD_CONTROL_TYPE_SEGMENT}, + {"inpaint", SD_CONTROL_TYPE_REPAINT}, + {"outpaint", SD_CONTROL_TYPE_REPAINT}, + }; + for (const auto& entry : aliases) { + if (!strcmp(str, entry.alias)) { + return entry.type; + } + } + return SD_CONTROL_TYPE_COUNT; +} + const char* sd_vae_format_name(enum sd_vae_format_t format) { switch (format) { case SD_VAE_FORMAT_AUTO: @@ -561,6 +612,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { "ref_images_count: %d\n" "ref_image_args: %s\n" "control_strength: %.2f\n" + "control_type: %s\n" "photo maker: {style_strength = %.2f, id_images_count = %d, id_embed_path = %s}\n" "VAE tiling: %s (temporal=%s, extra_tiling_args=%s)\n" "circular_x: %s\n" @@ -579,6 +631,7 @@ char* sd_img_gen_params_to_str(const sd_img_gen_params_t* sd_img_gen_params) { sd_img_gen_params->ref_images_count, SAFE_STR(sd_img_gen_params->ref_image_args), sd_img_gen_params->control_strength, + sd_control_type_name(sd_img_gen_params->control_type), sd_img_gen_params->pm_params.style_strength, sd_img_gen_params->pm_params.id_images_count, SAFE_STR(sd_img_gen_params->pm_params.id_embed_path), From b1f5a51a1ed98d89fef3013d2a1151a50d4b77ed Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 00:04:37 +0800 Subject: [PATCH 24/57] fix: chunk MiniMax-H3 QKV projections within backend buffer limits (#2286) --- src/core/ggml_runner.cpp | 25 ++++++- src/core/vulkan/int8.cpp | 8 +-- src/model/common/ggml_block.hpp | 3 +- src/model/diffusion/minimax_h3.hpp | 103 +++++++++++++++++++++-------- 4 files changed, 105 insertions(+), 34 deletions(-) diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 8c2c5ba218..2fa14b1508 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -211,6 +211,22 @@ ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const { return nullptr; } +static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { + auto buft = ggml_backend_get_default_buffer_type(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; + if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { + return false; + } + for (auto src : node->src) { + if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { + return false; + } + } + return true; +} + void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { if (sched == nullptr || gf == nullptr) { return; @@ -235,7 +251,14 @@ void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { continue; } - if (ggml_backend_supports_op(current, node)) { + // Operation support can allow tensors larger than the backend's allocation limit. + // Check inputs too, since cross-backend copies must also fit a single allocation. + if (!backend_can_allocate_op(current, node)) { + auto cpu = workspace_.cpu_backend(); + if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, cpu); + } + } else if (ggml_backend_supports_op(current, node)) { ggml_backend_sched_set_tensor_backend(sched, node, current); } } diff --git a/src/core/vulkan/int8.cpp b/src/core/vulkan/int8.cpp index aad36900be..2bac51bb5a 100644 --- a/src/core/vulkan/int8.cpp +++ b/src/core/vulkan/int8.cpp @@ -715,10 +715,10 @@ namespace sd_vulkan_i8 { const Parameters pc = {n, rows, k, op.bias ? 1u : 0u}; auto pipeline = op.kind == SDW4Kind::W4A4 ? SD_VK_PIPELINE(mul_mat_w4a4_convrot, 6, Parameters) : SD_VK_PIPELINE(mul_mat_w4a8_convrot, 8, Parameters); uint32_t tile_cols = 64, tile_rows = 32; - const bool large_w4a4 = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared; - const bool large_w4a8 = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared; - const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 && - aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4; + const bool large_w4a4 = op.kind == SDW4Kind::W4A4 && rows >= 512 && d.compute.shared_memory >= w4a4_large_shared; + const bool large_w4a8 = op.kind == SDW4Kind::W4A8 && rows >= 512 && d.compute.shared_memory >= w4a8_large_shared; + const bool decode_w4a4 = op.kind == SDW4Kind::W4A4 && d.w4a4_cm2 && k % 256 == 0 && n % 64 == 0 && rows >= 128 && + aligned(op.weight, 16) && sizeof(sd_vulkan_shaders::mul_mat_w4a4_convrot_decode_cm2) > 4; const uint32_t constants[] = {64, large_w4a4 || large_w4a8 ? 128u : 64u, 64}; const char* force_dp4a = std::getenv("GGML_VK_W4_FORCE_DP4A"); if (!(force_dp4a && std::atoi(force_dp4a)) && aligned(op.input, 16) && diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 0653ae1fff..39172bc3f6 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -974,7 +974,8 @@ class RMSNorm : public UnaryBlock { w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); } x = ggml_rms_norm(ctx->ggml_ctx, x, eps); - x = ggml_mul_inplace(ctx->ggml_ctx, x, w); + // The affine multiply may run on a different backend after normalization falls back. + x = ggml_mul(ctx->ggml_ctx, x, w); return x; } }; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index b5aec4397d..9a40643aab 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -243,8 +244,68 @@ namespace MiniMaxH3 { blocks["out_proj"] = std::make_shared(inner, hidden_size, false); } - ggml_tensor* project_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) { - return std::dynamic_pointer_cast(blocks["qkv_proj"])->forward(ctx, x); + std::vector split_qkv(GGMLRunnerContext* ctx, ggml_tensor* projected) { + std::vector result; + for (int i = 0; i < 3; ++i) { + result.push_back(ggml_view_4d(ctx->ggml_ctx, projected, head_dim, heads, + projected->ne[1], projected->ne[2] * projected->ne[3], + head_dim * projected->nb[0], projected->nb[1], projected->nb[2], + i * heads * head_dim * projected->nb[0])); + } + return result; + } + + std::vector project_qkv(GGMLRunnerContext* ctx, + const std::vector& inputs, + const std::function& prepare_input = {}) { + auto projection = std::dynamic_pointer_cast(blocks["qkv_proj"]); + int64_t sequence = 0; + for (auto input : inputs) { + sequence += input->ne[1]; + } + int64_t chunk_size = sequence; + 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; + size_t token_bytes = ggml_row_size(GGML_TYPE_F32, 3 * heads * head_dim) * inputs[0]->ne[2] * inputs[0]->ne[3]; + chunk_size = static_cast(std::min(chunk_size, std::max(1, max_size / token_bytes))); + } + if (chunk_size == sequence) { + ggml_tensor* projected = nullptr; + for (size_t i = 0; i < inputs.size(); ++i) { + auto input = prepare_input ? prepare_input(inputs[i], i) : inputs[i]; + auto part = projection->forward(ctx, input); + projected = projected == nullptr ? part : ggml_concat(ctx->ggml_ctx, projected, part, 1); + } + return split_qkv(ctx, projected); + } + + // Split Q/K/V before concatenating tokens so the full packed QKV never exceeds the buffer limit. + std::vector chunks[3]; + for (size_t i = 0; i < inputs.size(); ++i) { + auto input = inputs[i]; + for (int64_t start = 0; start < input->ne[1]; start += chunk_size) { + auto part = ggml_ext_slice(ctx->ggml_ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false); + if (!ggml_is_contiguous(part)) { + part = ggml_cont(ctx->ggml_ctx, part); + } + // Keep normalization/modulation adjacent to the projection's input quantization. + if (prepare_input) { + part = prepare_input(part, i); + } + auto qkv = split_qkv(ctx, projection->forward(ctx, part)); + for (int j = 0; j < 3; ++j) { + chunks[j].push_back(ggml_cont(ctx->ggml_ctx, qkv[j])); + } + } + } + std::vector result; + for (auto& parts : chunks) { + result.push_back(ggml_ext_vec_concat(ctx->ggml_ctx, parts, 2)); + } + return result; } ggml_tensor* project_output(GGMLRunnerContext* ctx, ggml_tensor* x) { @@ -252,22 +313,15 @@ namespace MiniMaxH3 { } ggml_tensor* attend(GGMLRunnerContext* ctx, - ggml_tensor* projected, + const std::vector& qkv, 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); + auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); + auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); + auto q = qkv[0]; + auto k = qkv[1]; + auto v = qkv[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); @@ -290,7 +344,7 @@ namespace MiniMaxH3 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr) { - return project_output(ctx, attend(ctx, project_qkv(ctx, x), pe)); + return project_output(ctx, attend(ctx, project_qkv(ctx, {x}), pe)); } }; @@ -447,16 +501,9 @@ namespace MiniMaxH3 { 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 qkv = attn->project_qkv(ctx, inputs, [&](ggml_tensor* input, size_t i) { + return modulate(norm1->forward(ctx, input), rows[i], 0, 1); + }); auto attended = attn->attend(ctx, qkv, pe); std::vector projections; for (const auto& segment : segments) { From 9234ad822a376a4724638467736118ed61d68ee9 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 00:09:48 +0800 Subject: [PATCH 25/57] feat: add RelayVSR video upscaling with CUDA window attention (#2285) --- README.md | 1 + docs/relay_vsr.md | 70 +++++++ examples/cli/main.cpp | 67 ++++++- examples/common/common.cpp | 1 + examples/common/common.h | 3 +- include/stable-diffusion.h | 31 +++ src/core/cuda/cute_attention.cu | 11 +- src/core/cuda/registration.cu | 2 + src/core/cuda/window_attention.cuh | 284 ++++++++++++++++++++++++++++ src/core/ggml_extensions.cpp | 59 ++++++ src/core/ggml_extensions.h | 12 ++ src/core/tensor.hpp | 42 +++-- src/core/util.cpp | 12 +- src/model/common/vsr.hpp | 221 ++++++++++++++++++++++ src/model/diffusion/relay_vsr.hpp | 293 +++++++++++++++++++++++++++++ src/model/upscaler/flash_vsr.hpp | 253 +++++++++++++++++++++++++ src/name_conversion.cpp | 33 ++++ src/name_conversion.h | 3 +- src/pipeline/relay_vsr.cpp | 209 ++++++++++++++++++++ src/pipeline/video_upscale.cpp | 143 ++++++++++++++ src/pipeline/video_upscale.h | 46 +++++ 21 files changed, 1764 insertions(+), 32 deletions(-) create mode 100644 docs/relay_vsr.md create mode 100644 src/core/cuda/window_attention.cuh create mode 100644 src/model/common/vsr.hpp create mode 100644 src/model/diffusion/relay_vsr.hpp create mode 100644 src/model/upscaler/flash_vsr.hpp create mode 100644 src/pipeline/relay_vsr.cpp create mode 100644 src/pipeline/video_upscale.cpp create mode 100644 src/pipeline/video_upscale.h diff --git a/README.md b/README.md index 28af00a697..2442865b2f 100644 --- a/README.md +++ b/README.md @@ -96,6 +96,7 @@ API and command-line option may change frequently.*** - Latent Consistency Models support (LCM/LCM-LoRA) - Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md) - Upscale images generated with [ESRGAN](./docs/esrgan.md) + - Upscale videos with [RelayVSR](./docs/relay_vsr.md) (experimental, 2x and 4x) - Supported backends - CPU (AVX, AVX2 and AVX512 support for x86 architectures) - CUDA diff --git a/docs/relay_vsr.md b/docs/relay_vsr.md new file mode 100644 index 0000000000..1352f41f19 --- /dev/null +++ b/docs/relay_vsr.md @@ -0,0 +1,70 @@ +# RelayVSR video super-resolution + +RelayVSR uses a Wan model to sample sparse reference latents, followed by the +Flash conditional decoder to reconstruct every input frame. The model code +follows [RelayVSR](https://github.com/kopperx/RelayVSR); the loader maps +[ComfyUI repackaged weights](https://huggingface.co/Comfy-Org/RelayVSR) to the +original module names. + +Use `-M vid_upscale` with an ordered directory of RGB images. Filenames are +sorted lexicographically. Supply the source frame rate with `--fps`. Output +keeps the input frame count. Use a video output supported by the CLI, such as +`.avi` or `.webm`, or a numbered PNG output such as `output/%04d.png`. + +The scale is an explicit, required argument. It is not read from checkpoint +metadata. Use matching Wan and Flash checkpoints: + +| Scale | `--diffusion-model` | `--vae` | +| --- | --- | --- | +| 2 | `relayvsr_x2_lightx2v_int8_convrot.safetensors` | `relayvsr_flash_vae_x2_fp32.safetensors` | +| 4 | `relayvsr_x4_int8_convrot.safetensors` | `relayvsr_flash_vae_x4_fp32.safetensors` | + +```sh +sd-cli -M vid_upscale \ + --diffusion-model models/diffusion_models/relayvsr_x2_lightx2v_int8_convrot.safetensors \ + --vae models/vae/relayvsr_flash_vae_x2_fp32.safetensors \ + --input-frames input_frames --upscale-factor 2 --fps 24 \ + --vsr-window-size 16 -s 42 -o output.avi +``` + +`--vsr-window-size` includes both reference endpoints. The default 16 samples +references at frame indices 0, 15, 30, and so on, always including the last +frame. Shared endpoints update Flash history in the following window but are +emitted only once. Single-frame input is also supported. + +Wan performs one flow-matching Euler step at timestep 1000 with guidance 1, +using the prompt embedding stored in the checkpoint. No text encoder or +ordinary Wan VAE is needed. Generic generation sampling settings do not +configure this fixed sampling path. + +The Flash scale cannot be verified from its tensor shapes: the x2 and x4 files +have identical shapes. The caller must select the correct pair. Wan's input +projection is checked against the requested scale. Inputs are replicate-padded +internally and outputs are cropped to exactly the input size times the scale. + +INT8 convrot uses the existing [quantized-weight execution path](int8_convrot.md). +Backend and parameter placement use `--backend` and `--params-backend`, with +Wan assigned to `diffusion` and Flash to `vae`. CUDA builds with CUTLASS use +fused window attention on compute capability 8.0 or newer. It loads spatial +neighborhoods directly, with FP16 Q/K/V and FP32 accumulation, and preserves +Wan's boundary-shifted windows and Flash's boundary-truncated windows. +Other backends use bounded spatial tiles with ordinary attention operators. +Set `GGML_CUDA_DISABLE_WINDOW_ATTENTION=1` to use that fallback on CUDA for +comparison. Floating-point results can differ between these implementations. + +Library callers initialize `sd_video_upscaler_ctx_params_t` with +`sd_video_upscaler_ctx_params_init()`, set the model paths, explicit scale and +backend options, and create a `video_upscaler_ctx_t` with +`new_video_upscaler_ctx()`. Model architecture is detected from +weight names and shapes; the scale remains caller-supplied. + +Initialize `sd_video_upscale_params_t` with `sd_video_upscale_params_init()`, +then set the RGB frames, frame count, FPS, seed and RelayVSR reference window. +Call `upscale_video()` and free the result with +`free_sd_video_generation_result()`. Defaults are 24 FPS, seed 42 and a +16-frame reference window. No sampler or scheduler selection is required. + +Each call starts fresh +temporal state. Model caches are bounded by the temporal windows; the public +API currently holds the complete input and output frame arrays in host memory. +Audio is not copied by this path. diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index 591a151123..bcf73595c6 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -33,10 +34,13 @@ const char* previews_str[] = { std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)"); struct SDCliParams { - SDMode mode = IMG_GEN; - std::string output_path = "output.png"; - int output_begin_idx = -1; - int compression_quality = 90; + std::string input_frames; + int video_upscale_factor = 0; + int vsr_window_size = 16; + SDMode mode = IMG_GEN; + std::string output_path = "output.png"; + int output_begin_idx = -1; + int compression_quality = 90; std::string image_path; std::string metadata_format = "text"; @@ -62,6 +66,7 @@ struct SDCliParams { ArgOptions options; options.string_options = { + {"", "--input-frames", "directory of ordered RGB frames for vid_upscale", 0, &input_frames}, {"-o", "--output", "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp", @@ -90,6 +95,8 @@ struct SDCliParams { }; options.int_options = { + {"", "--upscale-factor", "required video upscale factor (RelayVSR: 2 or 4)", &video_upscale_factor}, + {"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_size}, {"", "--preview-interval", "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1", @@ -190,7 +197,7 @@ struct SDCliParams { options.manual_options = { {"-M", "--mode", - "run mode, one of [img_gen, adetailer, vid_gen, upscale, convert, metadata], default: img_gen", + "run mode, one of [img_gen, adetailer, vid_gen, vid_upscale, upscale, convert, metadata], default: img_gen", on_mode_arg}, {"", "--preview", @@ -220,6 +227,10 @@ struct SDCliParams { } bool validate() { + if (mode == VID_UPSCALE && (input_frames.empty() || video_upscale_factor <= 0 || vsr_window_size < 2)) { + LOG_ERROR("vid_upscale requires --input-frames, positive --upscale-factor, and --vsr-window-size >= 2"); + return false; + } if (mode != METADATA) { if (output_path.length() == 0) { LOG_ERROR("error: the following arguments are required: output_path"); @@ -529,7 +540,7 @@ bool save_results(const SDCliParams& cli_params, return sucessful_reults != 0; } - if (cli_params.mode == VID_GEN && num_results > 1) { + if ((cli_params.mode == VID_GEN && num_results > 1) || cli_params.mode == VID_UPSCALE) { if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm") ext = ".avi"; std::string params = gen_params.embed_image_metadata @@ -664,6 +675,50 @@ int main(int argc, const char* argv[]) { return 0; } + if (cli_params.mode == VID_UPSCALE) { + if (ctx_params.diffusion_model_path.empty() || ctx_params.vae_path.empty() || gen_params.fps <= 0) { + LOG_ERROR("vid_upscale requires --diffusion-model, --vae and positive --fps"); + return 1; + } + std::vector owned_frames; + if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) { + LOG_ERROR("vid_upscale input contains no readable frames"); + return 1; + } + if (owned_frames.size() > size_t(INT_MAX)) + return 1; + std::vector frames; + for (const auto& frame : owned_frames) + frames.push_back(frame.get()); + sd_video_upscaler_ctx_params_t params; + sd_video_upscaler_ctx_params_init(¶ms); + params.diffusion_model_path = ctx_params.diffusion_model_path.c_str(); + params.vae_path = ctx_params.vae_path.c_str(); + params.upscale_factor = cli_params.video_upscale_factor; + params.n_threads = ctx_params.n_threads; + params.backend = ctx_params.effective_backend.c_str(); + params.params_backend = ctx_params.effective_params_backend.c_str(); + std::unique_ptr upscaler( + new_video_upscaler_ctx(¶ms), + free_video_upscaler_ctx); + if (!upscaler) + return 1; + sd_video_generation_result_t* output = nullptr; + sd_video_upscale_params_t upscale_params; + sd_video_upscale_params_init(&upscale_params); + upscale_params.frames = frames.data(); + upscale_params.frame_count = int(frames.size()); + upscale_params.fps = gen_params.fps; + upscale_params.window_size = cli_params.vsr_window_size; + upscale_params.seed = gen_params.seed; + if (!upscale_video(upscaler.get(), &upscale_params, &output)) + return 1; + std::string video_info = output->generation_info; + SDImageVec results; + results.adopt_video(output); + return save_results(cli_params, ctx_params, gen_params, results, int(results.size()), nullptr, video_info) ? 0 : 1; + } + if (!gen_params.ad_model_path.empty() && cli_params.mode != IMG_GEN && cli_params.mode != ADETAILER) { LOG_ERROR("--ad-model is only supported in image generation and adetailer modes"); return 1; diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 02c30880af..79349241cf 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -35,6 +35,7 @@ const char* const modes_str[] = { "convert", "upscale", "metadata", + "vid_upscale", }; static sd_vae_format_t str_to_vae_format(const std::string& value) { diff --git a/examples/common/common.h b/examples/common/common.h index 4ecbef4fa3..d4df4fad61 100644 --- a/examples/common/common.h +++ b/examples/common/common.h @@ -22,7 +22,7 @@ std::string argv_to_utf8(int index, const char** argv); #define BOOL_STR(b) ((b) ? "true" : "false") extern const char* const modes_str[]; -#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata" +#define SD_ALL_MODES_STR "img_gen, adetailer, vid_gen, convert, upscale, metadata, vid_upscale" enum SDMode { IMG_GEN, @@ -31,6 +31,7 @@ enum SDMode { CONVERT, UPSCALE, METADATA, + VID_UPSCALE, MODE_COUNT }; diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 1c61e4b150..a1d82e783a 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -615,6 +615,37 @@ SD_API int get_upscale_factor(upscaler_ctx_t* upscaler_ctx); // Reads model metadata only; returns 0 if the file is not a recognized RGB ESRGAN model. SD_API int get_upscaler_model_scale(const char* model_path); +typedef struct video_upscaler_ctx_t video_upscaler_ctx_t; + +typedef struct { + const char* diffusion_model_path; + const char* vae_path; + // Explicit output scale; RelayVSR requires 2 or 4 matching both checkpoints. + int upscale_factor; + int n_threads; + const char* backend; + const char* params_backend; +} sd_video_upscaler_ctx_params_t; + +typedef struct { + const sd_image_t* frames; + int frame_count; + int fps; + int64_t seed; + // RelayVSR reference window, including both endpoints; defaults to 16. + int window_size; +} sd_video_upscale_params_t; + +SD_API void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params); +SD_API void sd_video_upscale_params_init(sd_video_upscale_params_t* params); +SD_API video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params); +SD_API void free_video_upscaler_ctx(video_upscaler_ctx_t* ctx); +// Frames are RGB, have equal dimensions, and remain owned by the caller. +// Free the result with free_sd_video_generation_result(). +SD_API bool upscale_video(video_upscaler_ctx_t* ctx, + const sd_video_upscale_params_t* params, + sd_video_generation_result_t** result_out); + typedef struct adetailer_ctx_t adetailer_ctx_t; typedef struct { diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu index cdafc9983e..1eaefa4b01 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -150,7 +150,7 @@ namespace sd_fattn_cute { } } - template + template struct Softmax { decltype(make_tensor(Shape>{})) maximum, sum; template @@ -180,8 +180,9 @@ namespace sd_fattn_cute { } #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; + const float safe_max = Masked && next_max == -INFINITY ? 0.f : next_max; + const float value = exp2f(scores(m, n) * 1.4426950408889634f - safe_max * 1.4426950408889634f); + scores(m, n) = value; row_sum += value; } // Defer the four-lane sum reduction until the final output normalization. @@ -196,7 +197,7 @@ namespace sd_fattn_cute { float total = sum(m); total += __shfl_xor_sync(0xffffffff, total, 2); total += __shfl_xor_sync(0xffffffff, total, 1); - const float inverse = 1.f / total; + const float inverse = Masked && total == 0.f ? 0.f : 1.f / total; #pragma unroll for (int n = 0; n < size<1>(output); ++n) output(m, n) *= inverse; @@ -790,3 +791,5 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) { desc.priority = 400; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); } + +#include "window_attention.cuh" diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu index 1f48174313..d92be448ce 100644 --- a/src/core/cuda/registration.cu +++ b/src/core/cuda/registration.cu @@ -4,6 +4,7 @@ void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device); void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device); void sd_ggml_register_cuda_attention(ggml_backend_dev_t device); +void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device); #endif namespace SD_CUDA_NAMESPACE { @@ -39,6 +40,7 @@ namespace SD_CUDA_NAMESPACE { sd_ggml_register_cuda_i8_projection(device); sd_ggml_register_cuda_i8_gated_residual(device); sd_ggml_register_cuda_attention(device); + sd_ggml_register_cuda_window_attention(device); #endif } } diff --git a/src/core/cuda/window_attention.cuh b/src/core/cuda/window_attention.cuh new file mode 100644 index 0000000000..600cd91b64 --- /dev/null +++ b/src/core/cuda/window_attention.cuh @@ -0,0 +1,284 @@ +/* + * 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. + */ +#ifndef __SD_CORE_CUDA_WINDOW_ATTENTION_CUH__ +#define __SD_CORE_CUDA_WINDOW_ATTENTION_CUH__ + +namespace sd_fattn_cute { + struct WindowRegion { + int left, top, width, height; + }; + + static __device__ WindowRegion window_region(const SDWindowAttentionParams& p, int plane, int x, int y) { + int w = p.window_width[plane], h = p.window_height[plane]; + if (p.preserve_boundary) { + w = min(w, p.width); + h = min(h, p.height); + return {min(max(x - w / 2, 0), p.width - w), min(max(y - h / 2, 0), p.height - h), w, h}; + } + int left = max(x - w / 2, 0), top = max(y - h / 2, 0); + return {left, top, min(p.width, x + w / 2) - left, min(p.height, y + h / 2) - top}; + } + + template + __device__ __forceinline__ void window_copy(const cutlass::half_t* source, Dst& dst, const Id& id, int dim, Index index) { +#pragma unroll + for (int m = 0; m < size<1>(dst); ++m) { + const int row = index(get<0>(id(0, m, 0))); +#pragma unroll + for (int k = 0; k < size<2>(dst); ++k) { + if (row >= 0) { + const auto* src = source + int64_t(row) * dim + get<1>(id(0, m, k)); + const unsigned address = unsigned(__cvta_generic_to_shared(&dst(0, m, k))); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;" + : + : "r"(address), "l"(src)); + } else { + cute::clear(dst(_, m, k)); + } + } + } + } + + template + __global__ __launch_bounds__(128) void window_kernel(const cutlass::half_t* q, const cutlass::half_t* k, const cutlass::half_t* v, float* output, SDWindowAttentionParams p, int kv_heads) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + constexpr int BM = 64, BN = 64, TW = 8; + using T = Traits; + using Element = typename T::Element; + extern __shared__ char storage[]; + const int tid = threadIdx.x, head = blockIdx.z; + const int qx = blockIdx.x * TW, qy = blockIdx.y * (BM / TW); + const int tokens = p.width * p.height; + q += int64_t(head) * tokens * D; + const int64_t kv_offset = int64_t(head / (p.heads / kv_heads)) * tokens * p.planes * D; + k += kv_offset; + v += kv_offset; + 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 tQsQ = thread_copy.partition_D(sQ); + auto tKsK = thread_copy.partition_D(sK); + auto tVsV = thread_copy.partition_D(sV); + auto q_id = thread_copy.partition_S(make_identity_tensor(Shape, Int>{})); + auto kv_id = thread_copy.partition_S(make_identity_tensor(Shape, Int>{})); + auto query_index = [&](int row) { + const int x = qx + row % TW, y = qy + row / TW; + return x < p.width && y < p.height ? y * p.width + x : -1; + }; + window_copy(q, tQsQ, q_id, D, query_index); + cute::cp_async_fence(); + + 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>{})); + clear(acc_o); + Softmax<2 * size<1>(acc_o), true> softmax; + const float scale = rsqrtf(float(D)); + for (int plane = 0; plane < p.planes; ++plane) { + const auto first = window_region(p, plane, qx, qy); + const auto last = window_region(p, plane, min(qx + TW - 1, p.width - 1), min(qy + BM / TW - 1, p.height - 1)); + const int kw = last.left + last.width - first.left; + const int kh = last.top + last.height - first.top; + for (int start = 0; start < kw * kh; start += BN) { + auto key_index = [&](int row) { + row += start; + return row < kw * kh ? plane * tokens + (first.top + row / kw) * p.width + first.left + row % kw : -1; + }; + window_copy(k, tKsK, kv_id, D, key_index); + window_copy(v, tVsV, kv_id, D, key_index); + cute::cp_async_fence(); + cute::cp_async_wait<0>(); + __syncthreads(); + auto acc_s = partition_fragment_C(mma, Shape, Int>{}); + clear(acc_s); + gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK); +#pragma unroll + for (int i = 0; i < size(acc_s); ++i) { + const int qr = get<0>(score_id(i)), kr = start + get<1>(score_id(i)); + const int x = qx + qr % TW, y = qy + qr / TW; + const auto region = window_region(p, plane, x, y); + const int kx = first.left + kr % kw, ky = first.top + kr / kw; + const bool valid = x < p.width && y < p.height && kr < kw * kh && + kx >= region.left && kx < region.left + region.width && + ky >= region.top && ky < region.top + region.height; + acc_s(i) = valid ? acc_s(i) * scale : -INFINITY; + } + if (plane == 0 && start == 0) + softmax.template update(acc_s, acc_o); + else + 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); + __syncthreads(); + } + } + softmax.normalize(acc_o); + auto output_id = thread_mma.partition_C(make_identity_tensor(Shape, Int>{})); +#pragma unroll + for (int i = 0; i < size(acc_o); ++i) { + const int token = query_index(get<0>(output_id(i))); + if (token >= 0) + output[(int64_t(token) * p.heads + head) * D + get<1>(output_id(i))] = acc_o(i); + } +#else + // This translation unit also targets older GPUs; runtime dispatch excludes this kernel there. + asm volatile("trap;"); +#endif + } +} + +static bool window_attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) { + static const bool disabled = [] { + const char* value = std::getenv("GGML_CUDA_DISABLE_WINDOW_ATTENTION"); + return value && std::atoi(value) != 0; + }(); + SDWindowAttentionParams params; + const auto* info = cuda_device(device); + return !disabled && info && info->properties.major >= 8 && sd_ggml_window_attn_describe(node, ¶ms); +} + +static ggml_status window_attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) { + *workspace = 0; + for (int i = 0; i < 3; ++i) + *workspace += (size_t(ggml_nelements(nodes[0]->src[i])) * sizeof(half) + 255) / 256 * 256; + return GGML_STATUS_SUCCESS; +} + +static void window_attention_destroy(void*, void*) {} + +template +static cudaError_t window_attention_launch(const ggml_ext_cuda_context& context, const half* const* inputs, float* output, SDWindowAttentionParams params, int kv_heads) { + auto kernel = sd_fattn_cute::window_kernel; + 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)>>>( + reinterpret_cast(inputs[0]), reinterpret_cast(inputs[1]), + reinterpret_cast(inputs[2]), output, params, kv_heads); + return cudaGetLastError(); +} + +static ggml_status window_attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) { + 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; + SDWindowAttentionParams params; + auto* node = nodes[0]; + if (!sd_ggml_window_attn_describe(node, ¶ms)) + return GGML_STATUS_FAILED; + auto stream = static_cast(context.stream); + auto* base = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + const int dim = int(node->ne[0] / params.heads); + const half* inputs[3]; + for (int i = 0; i < 3; ++i) { + const auto* src = node->src[i]; + auto* packed = reinterpret_cast(base); + const int64_t heads = src->ne[0] / dim; + copy_to_half<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast(src->data), packed, + ggml_nelements(src), dim, src->ne[1], heads, src->nb[1], dim * sizeof(float), ggml_nbytes(src), 1.f); + if (cudaGetLastError() != cudaSuccess) + return GGML_STATUS_FAILED; + inputs[i] = packed; + base += (size_t(ggml_nelements(src)) * sizeof(half) + 255) / 256 * 256; + } + const int kv_heads = int(node->src[1]->ne[0] / dim); + auto* output = static_cast(node->data); + const auto error = dim == 64 ? window_attention_launch<64>(context, inputs, output, params, kv_heads) + : window_attention_launch<128>(context, inputs, output, params, kv_heads); + return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED; +} + +void sd_ggml_register_cuda_window_attention(ggml_backend_dev_t device) { + const auto* info = cuda_device(device); + if (!info || info->properties.major < 8) + return; + int previous; + if (cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info->ordinal) != cudaSuccess) + return; + cudaFuncAttributes attributes64{}, attributes128{}; + const auto error64 = cudaFuncGetAttributes(&attributes64, sd_fattn_cute::window_kernel<64>); + const auto error128 = cudaFuncGetAttributes(&attributes128, sd_fattn_cute::window_kernel<128>); + const auto restored = cudaSetDevice(previous); + // JIT can produce an SM80+ binary from older PTX containing only the trap. + if (restored != cudaSuccess || error64 != cudaSuccess || error128 != cudaSuccess || + std::min(attributes64.binaryVersion, attributes64.ptxVersion) < 80 || + std::min(attributes128.binaryVersion, attributes128.ptxVersion) < 80) + return; + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.window_attention.cuda.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_window_attn_op(); + desc.priority = 100; + desc.supports = window_attention_supports; + desc.prepare = window_attention_prepare; + desc.destroy = window_attention_destroy; + desc.execute = window_attention_execute; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); +} + +#endif diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index 25fa22265f..dfe5bca402 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -501,6 +501,65 @@ ggml_tensor* ggml_sol_attn(ggml_context* ctx, ggml_tensor* q, ggml_tensor* k, gg return ggml_ext_op_create(ctx, sd_ggml_sol_attn_op(), inputs, 3, ¶ms); } +static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { + if (count != 3 || !attributes) + return false; + SDWindowAttentionParams p; + std::memcpy(&p, attributes, sizeof(p)); + if (p.width <= 0 || p.height <= 0 || int64_t(p.width) * p.height > INT32_MAX / 4 || + p.heads <= 0 || p.planes < 1 || p.planes > 4 || (p.preserve_boundary != 0 && p.preserve_boundary != 1)) + return false; + for (int i = 0; i < p.planes; ++i) + if (p.window_height[i] <= 0 || p.window_width[i] <= 0 || + (!p.preserve_boundary && (p.window_height[i] % 2 || p.window_width[i] % 2))) + return false; + for (int i = 0; i < 3; ++i) { + const auto* t = inputs[i]; + if (!t || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || t->ne[2] != 1 || t->ne[3] != 1 || + t->ne[1] != int64_t(p.width) * p.height * (i ? p.planes : 1)) + return false; + } + const auto* q = inputs[0]; + const auto* k = inputs[1]; + if (q->ne[0] % p.heads || (q->ne[0] / p.heads != 64 && q->ne[0] / p.heads != 128) || + k->ne[0] <= 0 || k->ne[0] > q->ne[0] || q->ne[0] % k->ne[0] || + k->ne[0] % (q->ne[0] / p.heads) || !ggml_are_same_shape(k, inputs[2])) + return false; + *out = {GGML_TYPE_F32, {q->ne[0], q->ne[1], 1, 1}}; + return true; +} + +ggml_ext_op_t sd_ggml_window_attn_op() { + static const auto op = [] { + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), window_attn_infer}; + auto result = ggml_ext_register_op(&desc); + GGML_ASSERT(result); + return result; + }(); + return op; +} + +bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params) { + if (!node || ggml_ext_op_get(node) != sd_ggml_window_attn_op() || node->src[3] || !ggml_is_contiguous(node)) + return false; + const auto* attributes = ggml_ext_op_attributes(node); + ggml_ext_tensor_desc out; + if (!window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne))) + return false; + std::memcpy(params, attributes, sizeof(*params)); + return true; +} + +ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params) { + ggml_tensor* inputs[] = {q, k, v}; + ggml_ext_tensor_desc out; + if (!backend || !window_attn_infer(inputs, 3, ¶ms, &out)) + return nullptr; + sd_ggml_register_extensions(backend); + auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, ¶ms); + return ggml_backend_supports_op(backend, result) ? result : nullptr; +} + void sd_ggml_register_extensions(ggml_backend_t backend) { if (!backend) { return; diff --git a/src/core/ggml_extensions.h b/src/core/ggml_extensions.h index bc9ee9a5c4..87a51a7e5e 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -80,6 +80,18 @@ struct SDSolAttentionParams { float tau; }; +struct SDWindowAttentionParams { + int32_t width, height, heads, planes; + int16_t window_height[4], window_width[4]; + int32_t preserve_boundary; +}; + +ggml_ext_op_t sd_ggml_window_attn_op(); +bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params); +// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes]. +// Returns Q's layout, or nullptr when the backend has no compatible kernel. +ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params); + 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]. diff --git a/src/core/tensor.hpp b/src/core/tensor.hpp index a925dbffb3..715068c364 100644 --- a/src/core/tensor.hpp +++ b/src/core/tensor.hpp @@ -222,14 +222,28 @@ namespace sd { return result; } - for (int64_t flat = 0; flat < result.numel(); ++flat) { - std::vector out_coord = tensor_unravel_index(flat, out_shape); - std::vector src_coord(static_cast(dim()), 0); - for (size_t i = 0; i < dims.size(); ++i) { - src_coord[dims[i]] = out_coord[i]; + std::vector source_strides(dims.size(), 1), strides(dims.size()); + for (size_t i = 1; i < dims.size(); ++i) + source_strides[i] = source_strides[i - 1] * shape_[i - 1]; + for (size_t i = 0; i < dims.size(); ++i) + strides[i] = source_strides[dims[i]]; + parallel_for(0, result.numel(), 65536, [&](int64_t begin, int64_t end) { + auto coord = tensor_unravel_index(begin, out_shape); + int64_t source = 0; + for (size_t i = 0; i < dims.size(); ++i) + source += coord[i] * strides[i]; + for (int64_t flat = begin; flat < end; ++flat) { + result.data()[flat] = data_[source]; + for (size_t i = 0; i < dims.size(); ++i) { + ++coord[i]; + source += strides[i]; + if (coord[i] < out_shape[i]) + break; + source -= coord[i] * strides[i]; + coord[i] = 0; + } } - result[flat] = index(src_coord); - } + }); return result; } @@ -1037,11 +1051,13 @@ namespace sd { auto x_contributors = make_interpolate_contributors(input_width, output_width, mode, antialias); auto y_contributors = make_interpolate_contributors(input_height, output_height, mode, antialias); - for (int64_t plane = 0; plane < plane_count; ++plane) { - const int64_t input_plane_offset = plane * input_plane; - const int64_t output_plane_offset = plane * output_plane; - for (int64_t y = 0; y < output_height; ++y) { - const auto& y_axis = y_contributors[static_cast(y)]; + parallel_for(0, plane_count * output_height, std::max(1, 16384 / output_width), [&](int64_t begin, int64_t end) { + for (int64_t row = begin; row < end; ++row) { + const int64_t plane = row / output_height; + const int64_t y = row % output_height; + const int64_t input_plane_offset = plane * input_plane; + const int64_t output_plane_offset = plane * output_plane; + const auto& y_axis = y_contributors[static_cast(y)]; for (int64_t x = 0; x < output_width; ++x) { const auto& x_axis = x_contributors[static_cast(x)]; double value = 0.0; @@ -1055,7 +1071,7 @@ namespace sd { output.data()[output_plane_offset + y * output_width + x] = static_cast(value); } } - } + }); return output; } diff --git a/src/core/util.cpp b/src/core/util.cpp index a691b423c0..9de1ace24f 100644 --- a/src/core/util.cpp +++ b/src/core/util.cpp @@ -807,13 +807,11 @@ sd::Tensor sd_image_to_tensor(sd_image_t image, static_cast(image.height), static_cast(image.channel), 1}); - for (uint32_t iw = 0; iw < image.width; ++iw) { - for (uint32_t ih = 0; ih < image.height; ++ih) { - for (uint32_t ic = 0; ic < image.channel; ++ic) { - tensor.index(iw, ih, ic, 0) = sd_image_get_f32(image, iw, ih, ic, scale); - } - } - } + const int64_t pixels = int64_t(image.width) * image.height; + sd::tensor_for_each(tensor.numel(), [&](int64_t i) { + float value = image.data[(i % pixels) * image.channel + i / pixels]; + tensor.data()[i] = scale ? value / 255.f : value; + }); if (target_width >= 0 && target_height >= 0 && (tensor.shape()[0] != target_width || tensor.shape()[1] != target_height)) { tensor = sd::ops::interpolate(tensor, diff --git a/src/model/common/vsr.hpp b/src/model/common/vsr.hpp new file mode 100644 index 0000000000..c203dff10a --- /dev/null +++ b/src/model/common/vsr.hpp @@ -0,0 +1,221 @@ +#ifndef __SD_MODEL_COMMON_VSR_HPP__ +#define __SD_MODEL_COMMON_VSR_HPP__ + +#include +#include +#include +#include +#include +#include "model/common/ggml_block.hpp" +#include "model/common/rope.hpp" + +namespace VSR { + + struct GraphInputs { + std::map, sd::Tensor> rotations; + mutable std::unordered_map bindings; + + void begin_graph() { + bindings.clear(); + } + + void clear() { + bindings.clear(); + rotations.clear(); + } + + template + ggml_tensor* input(GGMLRunnerContext* ctx, const sd::Tensor& value) const { + auto existing = bindings.find(value.data()); + if (existing != bindings.end()) + return existing->second; + auto result = sd::make_ggml_tensor(ctx->ggml_ctx, value, false); + ctx->set_backend_tensor_data(result, value.data()); + bindings.emplace(value.data(), result); + return result; + } + + ggml_tensor* rope(GGMLRunnerContext* ctx, int width, int height, const std::vector& times, const std::vector& axes) { + std::vector key = {width, height, axes[0], axes[1], axes[2]}; + key.insert(key.end(), times.begin(), times.end()); + auto existing = rotations.find(key); + if (existing != rotations.end()) + return input(ctx, existing->second); + const int dim = axes[0] + axes[1] + axes[2]; + auto& pe = rotations.emplace(std::move(key), sd::Tensor({2, 2, dim / 2, int64_t(width) * height * int64_t(times.size())})).first->second; + float* out = pe.data(); + for (int t : times) { + for (int y = 0; y < height; ++y) { + for (int x = 0; x < width; ++x) { + const int pos[] = {t, y, x}; + for (int a = 0; a < 3; ++a) { + for (int i = 0; i < axes[a]; i += 2) { + const double angle = pos[a] * std::pow(10000.0, -double(i) / axes[a]); + *out++ = float(std::cos(angle)); + *out++ = float(-std::sin(angle)); + *out++ = float(std::sin(angle)); + *out++ = float(std::cos(angle)); + } + } + } + } + } + return input(ctx, pe); + } + }; + + static ggml_tensor* tokens(ggml_context* ctx, ggml_tensor* x) { + x = ggml_reshape_2d(ctx, x, x->ne[0] * x->ne[1], x->ne[2] * x->ne[3]); + return ggml_cont(ctx, ggml_transpose(ctx, x)); + } + + static ggml_tensor* image(ggml_context* ctx, ggml_tensor* x, int width, int height) { + const auto channels = x->ne[0]; + x = ggml_cont(ctx, ggml_transpose(ctx, x)); + return ggml_reshape_4d(ctx, x, width, height, channels, 1); + } + + static ggml_tensor* pixel_shuffle(ggml_context* ctx, ggml_tensor* x, int width, int height, int scale, bool channel_first = true) { + const int64_t channels = x->ne[0] / (scale * scale); + if (!channel_first) { + x = ggml_reshape_4d(ctx, x, channels, scale, scale, width * height); + x = ggml_cont(ctx, ggml_permute(ctx, x, 2, 0, 1, 3)); + } + x = ggml_reshape_4d(ctx, x, scale, scale, channels, width * height); + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 3, 1)); + x = ggml_reshape_4d(ctx, x, scale * width, height, scale, channels); + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); + return ggml_reshape_4d(ctx, x, scale * width, scale * height, channels, 1); + } + + struct AttentionTile { + int width; + int height; + sd::Tensor queries; + sd::Tensor keys; + sd::Tensor mask; + }; + + struct LocalAttention { + int width; + int height; + std::vector> windows; + bool preserve_boundary; + mutable std::vector tiles; + GraphInputs inputs; + + void begin_graph() { + inputs.begin_graph(); + } + + LocalAttention(int width, int height, const std::vector>& windows, bool preserve_boundary) + : width(width), height(height), windows(windows), preserve_boundary(preserve_boundary) {} + + void prepare_tiles() const { + if (!tiles.empty()) + return; + constexpr int tile_size = 16; + for (int y = 0; y < height; y += tile_size) { + for (int x = 0; x < width; x += tile_size) { + AttentionTile tile; + tile.width = std::min(tile_size, width - x); + tile.height = std::min(tile_size, height - y); + tile.queries = sd::Tensor({tile.width * tile.height}); + std::vector keys; + auto bounds = [&](int qy, int qx, int plane) { + auto [wh, ww] = windows[plane]; + if (preserve_boundary) { + wh = std::min(wh, height); + ww = std::min(ww, width); + const int top = std::clamp(qy - wh / 2, 0, height - wh); + const int left = std::clamp(qx - ww / 2, 0, width - ww); + return std::array{top, left, top + wh, left + ww}; + } + return std::array{std::max(0, qy - wh / 2), std::max(0, qx - ww / 2), + std::min(height, qy + wh / 2), std::min(width, qx + ww / 2)}; + }; + for (int p = 0; p < int(windows.size()); ++p) { + auto begin = bounds(y, x, p); + auto end = bounds(y + tile.height - 1, x + tile.width - 1, p); + for (int ky = begin[0]; ky < end[2]; ++ky) + for (int kx = begin[1]; kx < end[3]; ++kx) + keys.push_back(p * height * width + ky * width + kx); + } + tile.keys = sd::Tensor({int64_t(keys.size())}); + std::copy(keys.begin(), keys.end(), tile.keys.data()); + tile.mask = sd::Tensor({int64_t(keys.size()), tile.width * tile.height}); + for (int q = 0; q < tile.width * tile.height; ++q) { + int qy = y + q / tile.width; + int qx = x + q % tile.width; + tile.queries.data()[q] = qy * width + qx; + for (size_t k = 0; k < keys.size(); ++k) { + int p = keys[k] / (width * height); + int ky = (keys[k] % (width * height)) / width; + int kx = keys[k] % width; + auto b = bounds(qy, qx, p); + tile.mask.data()[q * keys.size() + k] = ky >= b[0] && ky < b[2] && kx >= b[1] && kx < b[3] + ? 0.f + : -std::numeric_limits::infinity(); + } + } + tiles.push_back(std::move(tile)); + } + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, int heads) const { + auto c = ctx->ggml_ctx; +#ifndef SD_USE_UPSTREAM_GGML + if (ctx->flash_attn_enabled && windows.size() <= 4) { + SDWindowAttentionParams params = {}; + params.width = width; + params.height = height; + params.heads = heads; + params.planes = int(windows.size()); + params.preserve_boundary = preserve_boundary; + for (size_t i = 0; i < windows.size(); ++i) { + params.window_height[i] = windows[i].first; + params.window_width[i] = windows[i].second; + } + if (auto out = sd_ggml_window_attn(c, ctx->backend, q, k, v, params)) + return out; + } +#endif + prepare_tiles(); + ggml_tensor* output = nullptr; + ggml_tensor* row = nullptr; + int row_width = 0; + for (const auto& tile : tiles) { + auto qi = inputs.input(ctx, tile.queries); + auto ki = inputs.input(ctx, tile.keys); + auto tq = ggml_get_rows(c, q, qi); + auto tk = ggml_get_rows(c, k, ki); + auto tv = ggml_get_rows(c, v, ki); + auto mask = inputs.input(ctx, tile.mask); + auto out = ggml_ext_attention_ext(c, ctx->backend, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled); + out = ggml_reshape_3d(c, out, out->ne[0], tile.width, tile.height); + row = row ? ggml_concat(c, row, out, 1) : out; + row_width += tile.width; + if (row_width == width) { + output = output ? ggml_concat(c, output, row, 2) : row; + row = nullptr; + row_width = 0; + } + } + return ggml_reshape_2d(c, output, output->ne[0], width * height); + } + }; + + static ggml_tensor* apply_rope(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, int heads) { + auto c = ctx->ggml_ctx; + const int64_t dim = x->ne[0]; + const int64_t count = x->ne[1]; + x = ggml_reshape_4d(c, x, dim / heads, heads, count, 1); + x = Rope::apply_rope(ctx, x, pe); + x = ggml_cont(c, ggml_permute(c, x, 0, 2, 1, 3)); + return ggml_reshape_2d(c, x, dim, count); + } + +} // namespace VSR + +#endif // __SD_MODEL_COMMON_VSR_HPP__ diff --git a/src/model/diffusion/relay_vsr.hpp b/src/model/diffusion/relay_vsr.hpp new file mode 100644 index 0000000000..a390aa34f1 --- /dev/null +++ b/src/model/diffusion/relay_vsr.hpp @@ -0,0 +1,293 @@ +#ifndef __SD_MODEL_DIFFUSION_RELAY_VSR_HPP__ +#define __SD_MODEL_DIFFUSION_RELAY_VSR_HPP__ + +#include "model/common/vsr.hpp" + +namespace RelayVSR { + + struct WanVSRConfig { + int dim = 3072; + int ffn_dim = 14336; + int num_heads = 24; + int num_layers = 30; + int in_channels = 48; + int text_dim = 4096; + int text_len = 512; + int hidden_dim1 = 768; + int hidden_dim2 = 3072; + int upscale_factor = 4; + + static WanVSRConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix, int scale) { + WanVSRConfig config; + config.upscale_factor = scale; + auto weight = [&](const std::string& name) -> const TensorStorage& { return weights.at(prefix + "." + name); }; + config.dim = int(weight("condition_embedder.time_embedder.linear_1.weight").ne[1]); + config.ffn_dim = int(weight("blocks.0.ffn.net.0.proj.weight").ne[1]); + config.num_heads = config.dim / 128; + config.text_dim = int(weight("prompt_embeds").ne[0]); + config.text_len = int(weight("prompt_embeds").ne[1]); + config.hidden_dim1 = int(weight("lq_proj.conv1.weight").ne[3]); + config.hidden_dim2 = int(weight("lq_proj.conv2.weight").ne[3]); + config.num_layers = 0; + while (weights.count(prefix + ".blocks." + std::to_string(config.num_layers) + ".attn1.to_q.weight")) + ++config.num_layers; + if ((scale != 2 && scale != 4) || weight("lq_proj.conv1.weight").ne[2] != 3 * (4 * scale) * (4 * scale) || + config.dim != 3072 || config.num_layers != 30 || config.text_dim != 4096 || + weight("proj_out.weight").ne[1] != 192) { + throw std::runtime_error("RelayVSR Wan weights do not match the requested scale or architecture"); + } + LOG_VERBOSE("relayvsr: dim = %d, num_layers = %d, ffn_dim = %d, scale = %d", config.dim, config.num_layers, config.ffn_dim, scale); + return config; + } + }; + + class LQ4xProj : public GGMLBlock { + int hidden_dim1_; + int hidden_dim2_; + + public: + explicit LQ4xProj(const WanVSRConfig& config) + : hidden_dim1_(config.hidden_dim1), hidden_dim2_(config.hidden_dim2) { + const int patch = 4 * config.upscale_factor; + blocks["conv1"] = std::make_shared(3 * patch * patch, hidden_dim1_, std::pair{1, 1}); + blocks["norm1"] = std::make_shared(hidden_dim1_, 1e-6f); + blocks["conv2"] = std::make_shared(hidden_dim1_, hidden_dim2_, std::pair{3, 3}, std::pair{2, 2}, std::pair{1, 1}); + blocks["norm2"] = std::make_shared(hidden_dim2_, 1e-6f); + blocks["out_proj"] = std::make_shared(hidden_dim2_, config.dim); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* unshuffled) { + auto conv1 = std::dynamic_pointer_cast(blocks["conv1"]); + auto conv2 = std::dynamic_pointer_cast(blocks["conv2"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto out_proj = std::dynamic_pointer_cast(blocks["out_proj"]); + auto c = ctx->ggml_ctx; + auto x = conv1->forward(ctx, unshuffled); + int width = int(x->ne[0]), height = int(x->ne[1]); + x = ggml_silu(c, norm1->forward(ctx, VSR::tokens(c, x))); + x = conv2->forward(ctx, VSR::image(c, x, width, height)); + x = ggml_silu(c, norm2->forward(ctx, VSR::tokens(c, x))); + return out_proj->forward(ctx, x); + } + }; + + class WanAttention : public GGMLBlock { + int dim_; + int heads_; + + public: + explicit WanAttention(const WanVSRConfig& config) + : dim_(config.dim), heads_(config.num_heads) { + for (const auto& name : {"to_q", "to_k", "to_v", "to_out.0"}) + blocks[name] = std::make_shared(dim_, dim_); + blocks["norm_q"] = std::make_shared(dim_, 1e-6f); + blocks["norm_k"] = std::make_shared(dim_, 1e-6f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* context, ggml_tensor* query_rope, ggml_tensor* key_rope, const VSR::LocalAttention* local, const std::string& cache_name) { + auto to_q = std::dynamic_pointer_cast(blocks["to_q"]); + auto to_k = std::dynamic_pointer_cast(blocks["to_k"]); + auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); + auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); + auto norm_q = std::dynamic_pointer_cast(blocks["norm_q"]); + auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); + auto c = ctx->ggml_ctx; + auto q = norm_q->forward(ctx, to_q->forward(ctx, x)); + auto past_k = ctx->get_cache_tensor(cache_name + ".key"); + auto past_v = ctx->get_cache_tensor(cache_name + ".value"); + ggml_tensor* k = nullptr; + ggml_tensor* v = nullptr; + if (context && past_k) { + k = past_k; + v = past_v; + } else { + k = norm_k->forward(ctx, to_k->forward(ctx, context ? context : x)); + v = to_v->forward(ctx, context ? context : x); + if (!context && past_k) { + k = ggml_concat(c, past_k, k, 1); + v = ggml_concat(c, past_v, v, 1); + } + auto retained = [&](ggml_tensor* value) { + if (context || value->ne[1] <= x->ne[1] * 3) + return value; + // The next attention window adds its current frame to these three previous frames. + const int64_t keep = x->ne[1] * 3; + return ggml_cont(c, ggml_view_2d(c, value, value->ne[0], keep, value->nb[1], (value->ne[1] - keep) * value->nb[1])); + }; + ctx->cache_tensor(cache_name + ".key", retained(k)); + ctx->cache_tensor(cache_name + ".value", retained(v)); + } + ggml_tensor* out; + if (local) { + q = VSR::apply_rope(ctx, q, query_rope, heads_); + k = VSR::apply_rope(ctx, k, key_rope, heads_); + out = local->forward(ctx, q, k, v, heads_); + } else { + out = ggml_ext_attention_ext(c, ctx->backend, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled); + } + return to_out_0->forward(ctx, out); + } + }; + + class WanTransformerBlock : public GGMLBlock { + int dim_; + + void init_params(ggml_context* ctx, const String2TensorStorage& weights, const std::string prefix) override { + params["scale_shift_table"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, dim_, 6, 1); + } + + public: + explicit WanTransformerBlock(const WanVSRConfig& config) + : dim_(config.dim) { + blocks["norm1"] = std::make_shared(dim_, 1e-6f, false); + blocks["norm2"] = std::make_shared(dim_, 1e-6f, true); + blocks["norm3"] = std::make_shared(dim_, 1e-6f, false); + blocks["attn1"] = std::make_shared(config); + blocks["attn2"] = std::make_shared(config); + blocks["ffn.net.0.proj"] = std::make_shared(dim_, config.ffn_dim); + blocks["ffn.net.2"] = std::make_shared(config.ffn_dim, dim_); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* text, ggml_tensor* timestep_proj, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, const std::string& name) { + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto norm3 = std::dynamic_pointer_cast(blocks["norm3"]); + auto ffn_net_0_proj = std::dynamic_pointer_cast(blocks["ffn.net.0.proj"]); + auto ffn_net_2 = std::dynamic_pointer_cast(blocks["ffn.net.2"]); + auto c = ctx->ggml_ctx; + auto modulation = ggml_add(c, timestep_proj, params["scale_shift_table"]); + auto pieces = ggml_ext_chunk(c, modulation, 6, 1); + auto attn1 = std::dynamic_pointer_cast(blocks["attn1"]); + auto attn2 = std::dynamic_pointer_cast(blocks["attn2"]); + auto y = attn1->forward(ctx, ggml_ext_modulate(ctx, norm1->forward(ctx, x), pieces[0], pieces[1]), nullptr, qpe, kpe, &local, name + ".attn1"); + x = ggml_add(c, x, ggml_mul(c, y, pieces[2])); + x = ggml_add(c, x, attn2->forward(ctx, norm2->forward(ctx, x), text, nullptr, nullptr, nullptr, name + ".attn2")); + y = ggml_ext_modulate(ctx, norm3->forward(ctx, x), pieces[3], pieces[4]); + y = ffn_net_0_proj->forward(ctx, y); + y = ffn_net_2->forward(ctx, ggml_ext_gelu(c, y)); + return ggml_add(c, x, ggml_mul(c, y, pieces[5])); + } + }; + + class WanTransformer3DModel : public GGMLBlock { + WanVSRConfig config_; + + void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override { + params["scale_shift_table"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config_.dim, 2, 1); + params["prompt_embeds"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config_.text_dim, config_.text_len); + } + + public: + explicit WanTransformer3DModel(const WanVSRConfig& config) + : config_(config) { + blocks["patch_embedding"] = std::make_shared(config.in_channels, config.dim, std::tuple{1, 2, 2}, std::tuple{1, 2, 2}); + blocks["lq_proj"] = std::make_shared(config); + blocks["condition_embedder.time_embedder.linear_1"] = std::make_shared(256, config.dim); + blocks["condition_embedder.time_embedder.linear_2"] = std::make_shared(config.dim, config.dim); + blocks["condition_embedder.time_proj"] = std::make_shared(config.dim, config.dim * 6); + blocks["condition_embedder.text_embedder.linear_1"] = std::make_shared(config.text_dim, config.dim); + blocks["condition_embedder.text_embedder.linear_2"] = std::make_shared(config.dim, config.dim); + for (int i = 0; i < config.num_layers; ++i) + blocks["blocks." + std::to_string(i)] = std::make_shared(config); + blocks["norm_out"] = std::make_shared(config.dim, 1e-6f, false); + blocks["proj_out"] = std::make_shared(config.dim, config.in_channels * 4); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* lq, ggml_tensor* timestep, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local) { + auto patch_embedding = std::dynamic_pointer_cast(blocks["patch_embedding"]); + auto condition_embedder_time_embedder_linear_1 = std::dynamic_pointer_cast(blocks["condition_embedder.time_embedder.linear_1"]); + auto condition_embedder_time_embedder_linear_2 = std::dynamic_pointer_cast(blocks["condition_embedder.time_embedder.linear_2"]); + auto condition_embedder_time_proj = std::dynamic_pointer_cast(blocks["condition_embedder.time_proj"]); + auto condition_embedder_text_embedder_linear_1 = std::dynamic_pointer_cast(blocks["condition_embedder.text_embedder.linear_1"]); + auto condition_embedder_text_embedder_linear_2 = std::dynamic_pointer_cast(blocks["condition_embedder.text_embedder.linear_2"]); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + auto c = ctx->ggml_ctx; + const int width = int(latent->ne[0]) / 2, height = int(latent->ne[1]) / 2; + auto x = patch_embedding->forward(ctx, latent); + x = VSR::tokens(c, ggml_reshape_4d(c, x, width, height, config_.dim, 1)); + x = ggml_add(c, x, std::dynamic_pointer_cast(blocks["lq_proj"])->forward(ctx, lq)); + auto e = condition_embedder_time_embedder_linear_1->forward(ctx, timestep); + e = condition_embedder_time_embedder_linear_2->forward(ctx, ggml_silu(c, e)); + auto e0 = condition_embedder_time_proj->forward(ctx, ggml_silu(c, e)); + e0 = ggml_reshape_3d(c, e0, config_.dim, 6, 1); + auto text = condition_embedder_text_embedder_linear_1->forward(ctx, params["prompt_embeds"]); + text = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(c, text)); + for (int i = 0; i < config_.num_layers; ++i) { + auto name = "blocks." + std::to_string(i); + x = std::dynamic_pointer_cast(blocks[name])->forward(ctx, x, text, e0, qpe, kpe, local, name); + sd::ggml_graph_cut::mark_graph_cut(x, "relayvsr." + name, "hidden"); + } + e = ggml_reshape_3d(c, e, config_.dim, 1, 1); + auto modulation = ggml_add(c, params["scale_shift_table"], e); + auto pieces = ggml_ext_chunk(c, modulation, 2, 1); + x = norm_out->forward(ctx, x); + x = ggml_ext_modulate(ctx, x, pieces[0], pieces[1]); + x = proj_out->forward(ctx, x); + x = VSR::pixel_shuffle(c, x, width, height, 2, false); + return ggml_reshape_4d(c, x, width * 2, height * 2, 1, config_.in_channels); + } + }; + + struct WanVSRRunner : GGMLRunner { + WanVSRConfig config; + WanTransformer3DModel model; + VSR::GraphInputs inputs; + std::map> attention; + sd::Tensor timestep; + int frame = 0; + + WanVSRRunner(ggml_backend_t backend, const String2TensorStorage& weights, int scale, std::shared_ptr manager) + : GGMLRunner(backend, manager), config(WanVSRConfig::detect_from_weights(weights, "wan", scale)), model(config), timestep({256}) { + model.init(params_ctx, weights, "wan"); + for (int i = 0; i < 128; ++i) { + double angle = 1000.0 * std::pow(10000.0, -double(i) / 128); + timestep.data()[i] = float(std::cos(angle)); + timestep.data()[i + 128] = float(std::sin(angle)); + } + } + + std::string get_desc() override { return "RelayVSR Wan"; } + void get_param_tensors(std::map& tensors) { model.get_param_tensors(tensors, "wan"); } + void reset() { + runner_end(); + free_cache_ctx_and_buffer(); + attention.clear(); + inputs.clear(); + frame = 0; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& noise, const sd::Tensor& lq) { + int width = int(noise.shape()[0]) / 2, height = int(noise.shape()[1]) / 2; + int count = std::min(frame + 1, 4); + if (!attention.count(count)) + attention[count] = std::make_unique(width, height, std::vector>(count, {22, 40}), true); + auto graph = [&]() { + inputs.begin_graph(); + attention[count]->begin_graph(); + auto gf = new_graph_custom(MAX_GRAPH_SIZE); + auto ctx = get_context(gf); + const int origin = std::max(0, frame - 7); + std::vector times; + for (int i = frame - count + 1; i <= frame; ++i) + times.push_back(i - origin); + auto qpe = inputs.rope(&ctx, width, height, {frame - origin}, {44, 42, 42}); + auto kpe = inputs.rope(&ctx, width, height, times, {44, 42, 42}); + auto out = model.forward(&ctx, make_input(noise), make_input(lq), make_input(timestep), qpe, kpe, *attention[count]); + ggml_build_forward_expand(gf, out); + return gf; + }; + auto prediction = restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), 4); + if (prediction.empty()) + return {}; + ++frame; + for (int64_t i = 0; i < prediction.numel(); ++i) + prediction.data()[i] = noise.data()[i] - prediction.data()[i]; + return prediction; + } + }; + +} // namespace RelayVSR + +#endif // __SD_MODEL_DIFFUSION_RELAY_VSR_HPP__ diff --git a/src/model/upscaler/flash_vsr.hpp b/src/model/upscaler/flash_vsr.hpp new file mode 100644 index 0000000000..470bfc3065 --- /dev/null +++ b/src/model/upscaler/flash_vsr.hpp @@ -0,0 +1,253 @@ +#ifndef __SD_MODEL_UPSCALER_FLASH_VSR_HPP__ +#define __SD_MODEL_UPSCALER_FLASH_VSR_HPP__ + +#include "model/common/vsr.hpp" + +namespace RelayVSR { + + struct FlashVSRConfig { + int dim = 512; + int depth = 10; + int num_heads = 8; + int num_kv_groups = 2; + int hidden_dim = 1408; + int head_dim = 1024; + int upscale_factor = 4; + + static FlashVSRConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix, int scale) { + FlashVSRConfig config; + config.upscale_factor = scale; + auto weight = [&](const std::string& name) -> const TensorStorage& { return weights.at(prefix + "." + name); }; + config.dim = int(weight("backbone.blocks.0.attention.query.weight").ne[1]); + config.num_heads = config.dim / 64; + config.num_kv_groups = int(weight("backbone.blocks.0.attention.key.weight").ne[1]) / 64; + config.hidden_dim = int(weight("backbone.blocks.0.feed_forward.down.weight").logical_ne0()); + config.head_dim = int(weight("spatial_head.0.weight").ne[1]); + config.depth = 0; + while (weights.count(prefix + ".backbone.blocks." + std::to_string(config.depth) + ".attention.query.weight")) + ++config.depth; + if ((scale != 2 && scale != 4) || config.dim != 512 || config.depth != 10 || config.num_kv_groups != 2 || + weight("condition_encoder.projection.weight").ne[2] != 48 || weight("residual_projection.weight").ne[1] != 768) + throw std::runtime_error("Unsupported RelayVSR Flash architecture"); + LOG_VERBOSE("flashvsr: dim = %d, depth = %d, hidden_dim = %d, scale = %d", config.dim, config.depth, config.hidden_dim, scale); + return config; + } + }; + + class LRFrameEncoder : public GGMLBlock { + public: + explicit LRFrameEncoder(const FlashVSRConfig& config) { + int channels[] = {3, 64, 128, 128}; + for (int i = 0; i < 3; ++i) { + int stride = i < 2 ? 2 : 4 / config.upscale_factor; + blocks["layers." + std::to_string(i * 2)] = std::make_shared(channels[i], channels[i + 1], std::pair{3, 3}, std::pair{stride, stride}, std::pair{1, 1}); + } + blocks["layers.6"] = std::make_shared(128, config.dim, std::pair{1, 1}); + blocks["norm"] = std::make_shared(config.dim, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto layers_6 = std::dynamic_pointer_cast(blocks["layers.6"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + for (int i = 0; i < 3; ++i) + x = ggml_silu(ctx->ggml_ctx, std::dynamic_pointer_cast(blocks["layers." + std::to_string(i * 2)])->forward(ctx, x)); + x = layers_6->forward(ctx, x); + return norm->forward(ctx, VSR::tokens(ctx->ggml_ctx, x)); + } + }; + + class ConditionEncoder : public GGMLBlock { + public: + explicit ConditionEncoder(const FlashVSRConfig& config) { + blocks["projection"] = std::make_shared(48, config.dim, std::pair{1, 1}); + blocks["norm"] = std::make_shared(config.dim, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto projection = std::dynamic_pointer_cast(blocks["projection"]); + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + return norm->forward(ctx, VSR::tokens(ctx->ggml_ctx, projection->forward(ctx, x))); + } + }; + + class ConditionedGroupedQueryAttention : public GGMLBlock { + FlashVSRConfig config_; + + public: + explicit ConditionedGroupedQueryAttention(const FlashVSRConfig& config) + : config_(config) { + blocks["query"] = std::make_shared(config.dim, config.dim); + blocks["output"] = std::make_shared(config.dim, config.dim); + for (const auto& name : {"key", "value", "condition_key", "condition_value"}) + blocks[name] = std::make_shared(config.dim, 64 * config.num_kv_groups); + blocks["key_norm"] = std::make_shared(64, 1e-5f); + blocks["value_norm"] = std::make_shared(64, 1e-5f); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, const std::string& name, int condition_slot, int start_slot, int end_slot, bool history) { + auto query = std::dynamic_pointer_cast(blocks["query"]); + auto output = std::dynamic_pointer_cast(blocks["output"]); + auto key_norm = std::dynamic_pointer_cast(blocks["key_norm"]); + auto value_norm = std::dynamic_pointer_cast(blocks["value_norm"]); + auto c = ctx->ggml_ctx; + const bool condition = condition_slot >= 0; + auto key = std::dynamic_pointer_cast(blocks[condition ? "condition_key" : "key"]); + auto value = std::dynamic_pointer_cast(blocks[condition ? "condition_value" : "value"]); + auto q = query->forward(ctx, x); + auto k = key->forward(ctx, x); + auto v = value->forward(ctx, x); + const int64_t count = x->ne[1]; + k = key_norm->forward(ctx, ggml_reshape_3d(c, k, 64, config_.num_kv_groups, count)); + v = value_norm->forward(ctx, ggml_reshape_3d(c, v, 64, config_.num_kv_groups, count)); + k = ggml_reshape_2d(c, k, 64 * config_.num_kv_groups, count); + v = ggml_reshape_2d(c, v, 64 * config_.num_kv_groups, count); + auto cache_prefix = name + (condition ? ".condition." + std::to_string(condition_slot) : ".video"); + auto past_k = history ? ctx->get_cache_tensor(name + ".video.key") : nullptr; + auto past_v = history ? ctx->get_cache_tensor(name + ".video.value") : nullptr; + ctx->cache_tensor(cache_prefix + ".key", k); + ctx->cache_tensor(cache_prefix + ".value", v); + if (!condition) { + auto join = [&](ggml_tensor* current, ggml_tensor* past, const char* suffix) { + auto start = ctx->get_cache_tensor(name + ".condition." + std::to_string(start_slot) + suffix); + auto end = ctx->get_cache_tensor(name + ".condition." + std::to_string(end_slot) + suffix); + if (!start || !end) + throw std::runtime_error("RelayVSR Flash endpoint cache is missing"); + auto value = ggml_concat(c, start, end, 1); + if (past) + value = ggml_concat(c, value, past, 1); + return ggml_concat(c, value, current, 1); + }; + k = join(k, past_k, ".key"); + v = join(v, past_v, ".value"); + } + q = VSR::apply_rope(ctx, q, qpe, config_.num_heads); + k = VSR::apply_rope(ctx, k, kpe, config_.num_kv_groups); + return output->forward(ctx, local.forward(ctx, q, k, v, config_.num_heads)); + } + }; + + class ConditionedTransformerBlock : public GGMLBlock { + public: + explicit ConditionedTransformerBlock(const FlashVSRConfig& config) { + blocks["attention_norm"] = std::make_shared(config.dim, 1e-5f); + blocks["attention"] = std::make_shared(config); + blocks["feed_forward_norm"] = std::make_shared(config.dim, 1e-5f); + blocks["feed_forward.gate_up_proj"] = std::make_shared(config.dim, config.hidden_dim * 2, false); + blocks["feed_forward.down"] = std::make_shared(config.hidden_dim, config.dim, false); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, const std::string& name, int condition_slot, int start_slot, int end_slot, bool history) { + auto attention_norm = std::dynamic_pointer_cast(blocks["attention_norm"]); + auto feed_forward_norm = std::dynamic_pointer_cast(blocks["feed_forward_norm"]); + auto feed_forward_gate_up_proj = std::dynamic_pointer_cast(blocks["feed_forward.gate_up_proj"]); + auto feed_forward_down = std::dynamic_pointer_cast(blocks["feed_forward.down"]); + auto c = ctx->ggml_ctx; + auto attention = std::dynamic_pointer_cast(blocks["attention"]); + auto y = attention->forward(ctx, attention_norm->forward(ctx, x), qpe, kpe, local, name, condition_slot, start_slot, end_slot, history); + x = ggml_add(c, x, y); + y = feed_forward_gate_up_proj->forward(ctx, feed_forward_norm->forward(ctx, x)); + auto parts = ggml_ext_chunk(c, y, 2, 0); + y = ggml_mul(c, ggml_silu(c, parts[0]), parts[1]); + return ggml_add(c, x, feed_forward_down->forward(ctx, y)); + } + }; + + class FlashVSRConditionNet : public GGMLBlock { + FlashVSRConfig config_; + + void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override { + params["video_type_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config_.dim, 1, 1); + params["condition_type_embedding"] = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, config_.dim, 1, 1); + } + + public: + explicit FlashVSRConditionNet(const FlashVSRConfig& config) + : config_(config) { + blocks["lr_encoder"] = std::make_shared(config); + blocks["condition_encoder"] = std::make_shared(config); + for (int i = 0; i < config.depth; ++i) + blocks["backbone.blocks." + std::to_string(i)] = std::make_shared(config); + blocks["spatial_head.0"] = std::make_shared(config.dim, config.head_dim); + blocks["residual_projection"] = std::make_shared(config.head_dim, 768); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* qpe, ggml_tensor* kpe, const VSR::LocalAttention& local, int condition_slot, int start_slot, int end_slot, bool history) { + auto spatial_head_0 = std::dynamic_pointer_cast(blocks["spatial_head.0"]); + auto residual_projection = std::dynamic_pointer_cast(blocks["residual_projection"]); + auto c = ctx->ggml_ctx; + const bool condition = condition_slot >= 0; + ggml_tensor* x; + if (condition) + x = std::dynamic_pointer_cast(blocks["condition_encoder"])->forward(ctx, input); + else + x = std::dynamic_pointer_cast(blocks["lr_encoder"])->forward(ctx, input); + x = ggml_add(c, x, params[condition ? "condition_type_embedding" : "video_type_embedding"]); + for (int i = 0; i < config_.depth; ++i) { + auto name = "backbone.blocks." + std::to_string(i); + x = std::dynamic_pointer_cast(blocks[name])->forward(ctx, x, qpe, kpe, local, name, condition_slot, start_slot, end_slot, history); + sd::ggml_graph_cut::mark_graph_cut(x, "flashvsr." + name, "hidden"); + } + if (condition) + return x; + x = residual_projection->forward(ctx, ggml_silu(c, spatial_head_0->forward(ctx, x))); + return VSR::pixel_shuffle(c, x, local.width, local.height, 16); + } + }; + + struct FlashVSRRunner : GGMLRunner { + FlashVSRConfig config; + FlashVSRConditionNet model; + VSR::GraphInputs inputs; + std::map> attention; + bool history = false; + + FlashVSRRunner(ggml_backend_t backend, const String2TensorStorage& weights, int scale, std::shared_ptr manager) + : GGMLRunner(backend, manager), config(FlashVSRConfig::detect_from_weights(weights, "flash", scale)), model(config) { + model.init(params_ctx, weights, "flash"); + } + std::string get_desc() override { return "RelayVSR Flash"; } + void get_param_tensors(std::map& tensors) { model.get_param_tensors(tensors, "flash"); } + void reset() { + runner_end(); + free_cache_ctx_and_buffer(); + attention.clear(); + inputs.clear(); + history = false; + } + + sd::Tensor compute(int n_threads, const sd::Tensor& input, int condition_slot, int start_slot = 0, int end_slot = 1) { + const bool condition = condition_slot >= 0; + const int stride = condition ? 1 : 16 / config.upscale_factor; + const int width = int(input.shape()[0]) / stride, height = int(input.shape()[1]) / stride; + const int planes = condition ? 1 : history ? 4 + : 3; + if (!attention.count(planes)) { + std::vector> windows(planes, {24, 24}); + windows[0] = {48, 48}; + if (!condition) + windows[1] = {48, 48}; + attention[planes] = std::make_unique(width, height, windows, false); + } + auto graph = [&]() { + inputs.begin_graph(); + attention[planes]->begin_graph(); + auto gf = new_graph_custom(MAX_GRAPH_SIZE); + auto ctx = get_context(gf); + auto qpe = inputs.rope(&ctx, width, height, {0}, {16, 24, 24}); + std::vector times = condition ? std::vector{0} : history ? std::vector{0, 0, -1, 0} + : std::vector{0, 0, 0}; + auto kpe = inputs.rope(&ctx, width, height, times, {16, 24, 24}); + auto out = model.forward(&ctx, make_input(input), qpe, kpe, *attention[planes], condition_slot, start_slot, end_slot, history && !condition); + ggml_build_forward_expand(gf, out); + return gf; + }; + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(graph, n_threads, false), condition ? 2 : 4); + if (!condition && !result.empty()) + history = true; + return result; + } + }; + +} // namespace RelayVSR + +#endif // __SD_MODEL_UPSCALER_FLASH_VSR_HPP__ diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 1f19b11879..78a5e1061a 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -5,6 +5,39 @@ #include "core/util.h" #include "name_conversion.h" +std::string convert_relayvsr_tensor_name(std::string name) { + if (!starts_with(name, "wan.")) + return name; + const std::vector> mappings = { + {"positive_conditioning", "prompt_embeds"}, + {"lq_proj_in.linear.", "lq_proj.out_proj."}, + {"lq_proj_in.", "lq_proj."}, + {"text_embedding.0.", "condition_embedder.text_embedder.linear_1."}, + {"text_embedding.2.", "condition_embedder.text_embedder.linear_2."}, + {"time_embedding.0.", "condition_embedder.time_embedder.linear_1."}, + {"time_embedding.2.", "condition_embedder.time_embedder.linear_2."}, + {"time_projection.1.", "condition_embedder.time_proj."}, + {"head.head.", "proj_out."}, + {"head.modulation", "scale_shift_table"}, + {".norm3.", ".norm2."}, + {".modulation", ".scale_shift_table"}, + {".ffn.0.", ".ffn.net.0.proj."}, + {".ffn.2.", ".ffn.net.2."}, + {".self_attn.", ".attn1."}, + {".cross_attn.", ".attn2."}, + {".q.", ".to_q."}, + {".k.", ".to_k."}, + {".v.", ".to_v."}, + {".o.", ".to_out.0."}, + }; + for (const auto& mapping : mappings) { + auto pos = name.find(mapping.first); + if (pos != std::string::npos) + name.replace(pos, mapping.first.size(), mapping.second); + } + return name; +} + void replace_with_name_map(std::string& name, const std::vector>& name_map) { for (auto kv : name_map) { size_t pos = name.find(kv.first); diff --git a/src/name_conversion.h b/src/name_conversion.h index 3fefcf7808..8d90c7e2cb 100644 --- a/src/name_conversion.h +++ b/src/name_conversion.h @@ -10,5 +10,6 @@ bool is_diffusion_model_name(const std::string& name); bool is_first_stage_model_name(const std::string& name); std::string convert_tensor_name(std::string name, SDVersion version); +std::string convert_relayvsr_tensor_name(std::string name); -#endif // __NAME_CONVERSTION_H__ \ No newline at end of file +#endif // __NAME_CONVERSTION_H__ diff --git a/src/pipeline/relay_vsr.cpp b/src/pipeline/relay_vsr.cpp new file mode 100644 index 0000000000..bc5cad7189 --- /dev/null +++ b/src/pipeline/relay_vsr.cpp @@ -0,0 +1,209 @@ +#include "pipeline/video_upscale.h" + +#include + +#include "core/rng_philox.hpp" +#include "generation_info.h" +#include "model/diffusion/relay_vsr.hpp" +#include "model/upscaler/flash_vsr.hpp" +#include "name_conversion.h" + +namespace sd { + namespace pipeline { + + static sd::Tensor prepare_vsr_frame(const sd_image_t& frame, int width, int height) { + auto result = sd_image_to_tensor(frame, -1, -1, false); + result /= 127.5f; + result -= 1.f; + if (width > int(frame.width)) { + auto edge = sd::ops::slice(result, 0, frame.width - 1, frame.width); + auto padding = sd::zeros({width - frame.width, frame.height, 3, 1}) + edge; + result = sd::ops::concat(result, padding, 0); + } + if (height > int(frame.height)) { + auto edge = sd::ops::slice(result, 1, frame.height - 1, frame.height); + auto padding = sd::zeros({width, height - frame.height, 3, 1}) + edge; + result = sd::ops::concat(result, padding, 1); + } + return result; + } + + static sd::Tensor prepare_vsr_lq(const sd::Tensor& frame, int scale) { + auto input = frame; + if (scale == 4) + input = sd::ops::interpolate(frame, {frame.shape()[0] * 4, frame.shape()[1] * 4, 3, 1}, sd::ops::InterpolateMode::Bilinear); + const int patch = 4 * scale; + const int width = int(input.shape()[0]), height = int(input.shape()[1]); + input.reshape_({patch, width / patch, patch, height / patch, 3, 1}); + input.permute_({1, 3, 0, 2, 4, 5}); + input.reshape_({width / patch, height / patch, 3 * patch * patch, 1}); + return input; + } + + class RelayVSRUpscaler : public VideoUpscaleModel { + std::unique_ptr wan_; + std::unique_ptr flash_; + int scale_ = 0; + int n_threads_ = 1; + + public: + bool init(VideoUpscaleRuntime& runtime, const String2TensorStorage& source_weights, int upscale_factor) override { + if (upscale_factor != 2 && upscale_factor != 4) { + LOG_ERROR("RelayVSR requires an explicit upscale factor of 2 or 4 matching both checkpoints"); + return false; + } + scale_ = upscale_factor; + n_threads_ = runtime.n_threads; + auto wan_backend = runtime.backends.runtime_backend(SDBackendModule::DIFFUSION); + auto flash_backend = runtime.backends.runtime_backend(SDBackendModule::VAE); + auto wan_params = runtime.backends.params_backend(SDBackendModule::DIFFUSION); + auto flash_params = runtime.backends.params_backend(SDBackendModule::VAE); + String2TensorStorage weights; + std::map source_names; + for (const auto& [name, storage] : source_weights) { + auto canonical = name; + if (starts_with(canonical, "diffusion.")) + canonical = "wan." + canonical.substr(10); + else if (starts_with(canonical, "vae.")) + canonical = "flash." + canonical.substr(4); + canonical = convert_relayvsr_tensor_name(canonical); + if (!source_names.emplace(canonical, name).second) + throw std::runtime_error("Duplicate RelayVSR weight name: " + canonical); + weights[canonical] = storage; + weights[canonical].name = canonical; + } + wan_ = std::make_unique(wan_backend, weights, upscale_factor, runtime.manager); + flash_ = std::make_unique(flash_backend, weights, upscale_factor, runtime.manager); + auto register_model = [&](auto& runner, ModelComponent component, SDBackendModule module, ggml_backend_t compute, ggml_backend_t params) { + std::map tensors, sources; + runner->get_param_tensors(tensors); + for (const auto& [name, tensor] : tensors) { + auto source = source_names.find(name); + if (source == source_names.end()) + throw std::runtime_error("Missing RelayVSR tensor: " + name); + sources.emplace(source->second, tensor); + ggml_set_name(tensor, source->second.c_str()); + } + runner->set_flash_attention_enabled(true); + return runtime.manager->register_param_tensors(component, std::move(sources), + runtime.backends.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, + compute, params); + }; + if (!register_model(wan_, ModelComponent::Diffusion, SDBackendModule::DIFFUSION, wan_backend, wan_params) || + !register_model(flash_, ModelComponent::VAE, SDBackendModule::VAE, flash_backend, flash_params) || + !runtime.manager->validate_registered_tensors()) + return false; + return true; + } + + void end() override { + if (wan_) + wan_->runner_end(); + if (flash_) + flash_->runner_end(); + } + + bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t* result) override { + const auto* frames = params.frames; + const int frame_count = params.frame_count, fps = params.fps, window_size = params.window_size; + const int64_t seed = params.seed; + if (window_size < 2) { + LOG_ERROR("RelayVSR requires a reference window of at least two frames"); + return false; + } + if (frames[0].width > uint32_t(INT_MAX / scale_ - 32) || frames[0].height > uint32_t(INT_MAX / scale_ - 32)) + return false; + struct Reset { + RelayVSRUpscaler* model; + ~Reset() { + model->wan_->reset(); + model->flash_->reset(); + } + } reset{this}; + wan_->reset(); + flash_->reset(); + const int stride = 16 / scale_; + const int alignment = 2 * stride; + const int width = (int(frames[0].width) + alignment - 1) / alignment * alignment; + const int height = (int(frames[0].height) + alignment - 1) / alignment * alignment; + PhiloxRNG rng; + rng.manual_seed(uint64_t(seed)); + auto keyframe = [&](int index, int slot) { + LOG_INFO("RelayVSR reference frame %d/%d", index + 1, frame_count); + auto input = prepare_vsr_frame(frames[index], width, height); + sd::Tensor noise({width / stride, height / stride, 1, 48}); + auto values = rng.randn(uint32_t(noise.numel())); + std::copy(values.begin(), values.end(), noise.data()); + auto latent = wan_->compute(n_threads_, noise, prepare_vsr_lq(input, scale_)); + if (latent.empty()) + return false; + latent.reshape_({width / stride, height / stride, 48, 1}); + return !flash_->compute(n_threads_, latent, slot).empty(); + }; + int start_slot = 0; + if (!keyframe(0, start_slot)) + return false; + for (int first = 0; first < frame_count;) { + const int last = first + std::min(window_size - 1, frame_count - 1 - first); + int end_slot = last == first ? start_slot : 1 - start_slot; + if (last != first && !keyframe(last, end_slot)) + return false; + flash_->history = false; + for (int i = first; i <= last; ++i) { + auto input = prepare_vsr_frame(frames[i], width, height); + auto residual = flash_->compute(n_threads_, input, -1, start_slot, end_slot); + if (residual.empty()) + return false; + if (first != 0 && i == first) + continue; + auto output = sd::ops::interpolate(input, {width * scale_, height * scale_, 3, 1}, sd::ops::InterpolateMode::Bicubic); + auto& image = result->frames[i]; + image.width = frames[i].width * scale_; + image.height = frames[i].height * scale_; + image.channel = 3; + output += residual; + output += 1.f; + output *= 127.5f; + output = sd::ops::slice(sd::ops::slice(output, 0, 0, image.width), 1, 0, image.height); + if (!std::all_of(output.values().begin(), output.values().end(), [](float value) { return std::isfinite(value); })) + throw std::runtime_error("RelayVSR produced non-finite pixels"); + output = sd::ops::clamp(output, 0.f, 255.f); + output.permute_({2, 0, 1, 3}); + image.data = static_cast(malloc(size_t(image.width) * image.height * 3)); + if (!image.data) + throw std::bad_alloc(); + std::transform(output.values().begin(), output.values().end(), image.data, [](float value) { return uint8_t(value); }); + LOG_INFO("RelayVSR output frame %d/%d", i + 1, frame_count); + if (auto progress = sd_get_progress_callback()) + progress(i + 1, frame_count, 0.f, sd_get_progress_callback_data()); + } + if (last == frame_count - 1) + break; + first = last; + start_slot = end_slot; + } + auto info = sd::upscale_generation_info(scale_); + info["kind"] = "video"; + info["upscale"]["model"] = "relayvsr"; + info["upscale"]["window_size"] = window_size; + info["seed"] = seed; + info["fps"] = fps; + info["video_frames"] = frame_count; + result->generation_info = sd::copy_generation_info(info); + if (!result->generation_info) + throw std::bad_alloc(); + return true; + } + }; + + bool is_relay_vsr_model(const String2TensorStorage& weights) { + return (weights.count("diffusion.positive_conditioning") || weights.count("diffusion.prompt_embeds")) && + (weights.count("diffusion.lq_proj_in.conv1.weight") || weights.count("diffusion.lq_proj.conv1.weight")); + } + + std::unique_ptr create_relay_vsr_upscaler() { + return std::make_unique(); + } + + } // namespace pipeline +} // namespace sd diff --git a/src/pipeline/video_upscale.cpp b/src/pipeline/video_upscale.cpp new file mode 100644 index 0000000000..5b34a9288c --- /dev/null +++ b/src/pipeline/video_upscale.cpp @@ -0,0 +1,143 @@ +#include "pipeline/video_upscale.h" + +#include + +#include "core/util.h" + +namespace sd { + namespace pipeline { + + VideoUpscaler::~VideoUpscaler() { + if (model_) + model_->end(); + // ModelManager retains parameter pointers owned by the model runners. + runtime_.manager.reset(); + model_.reset(); + } + + bool VideoUpscaler::init(const sd_video_upscaler_ctx_params_t& params) { + if (!params.diffusion_model_path || !*params.diffusion_model_path || + !params.vae_path || !*params.vae_path || params.upscale_factor <= 0) { + LOG_ERROR("video upscale requires diffusion and decoder model paths and a positive upscale factor"); + return false; + } + upscale_factor_ = params.upscale_factor; + runtime_.n_threads = params.n_threads > 0 ? params.n_threads : sd_get_num_physical_cores(); + executor_ = std::make_unique(runtime_.n_threads); + sd::ParallelScope scope(executor_.get()); + std::string error; + if (!runtime_.backends.init(params.backend, params.params_backend, nullptr, &error)) { + LOG_ERROR("video upscaler backend: %s", error.c_str()); + return false; + } + for (auto module : {SDBackendModule::DIFFUSION, SDBackendModule::VAE}) { + if (!runtime_.backends.runtime_backend(module) || !runtime_.backends.params_backend(module)) + return false; + } + ModelLoader loader; + if (!loader.init_from_file(params.diffusion_model_path, "diffusion.") || + !loader.init_from_file(params.vae_path, "vae.")) + return false; + const auto& weights = loader.get_tensor_storage_map(); + if (is_relay_vsr_model(weights)) + model_ = create_relay_vsr_upscaler(); + else { + LOG_ERROR("unsupported video upscale model architecture"); + return false; + } + runtime_.manager = std::make_shared(); + runtime_.manager->set_n_threads(runtime_.n_threads); + if (!runtime_.manager->set_loader(loader)) + return false; + return model_->init(runtime_, weights, upscale_factor_); + } + + bool VideoUpscaler::upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t** result_out) { + if (!result_out) + return false; + *result_out = nullptr; + if (!model_ || !params.frames || params.frame_count <= 0 || params.fps <= 0) { + LOG_ERROR("video upscale requires nonempty RGB frames and positive FPS"); + return false; + } + const auto& first = params.frames[0]; + if (!first.width || !first.height || first.width > uint32_t(INT_MAX / upscale_factor_) || + first.height > uint32_t(INT_MAX / upscale_factor_)) + return false; + for (int i = 0; i < params.frame_count; ++i) { + const auto& frame = params.frames[i]; + if (!frame.data || frame.channel != 3 || frame.width != first.width || frame.height != first.height) { + LOG_ERROR("video upscale input frames must have equal dimensions and three RGB channels"); + return false; + } + } + sd::ParallelScope scope(executor_.get()); + auto result = std::unique_ptr( + static_cast(calloc(1, sizeof(sd_video_generation_result_t))), free_sd_video_generation_result); + if (!result) + return false; + result->frames = static_cast(calloc(params.frame_count, sizeof(sd_image_t))); + if (!result->frames) + return false; + result->frame_count = params.frame_count; + result->fps = params.fps; + if (!model_->upscale(params, result.get())) + return false; + *result_out = result.release(); + return true; + } + + } // namespace pipeline +} // namespace sd + +struct video_upscaler_ctx_t { + sd::pipeline::VideoUpscaler upscaler; +}; + +void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params) { + if (!params) + return; + *params = {}; + params->n_threads = sd_get_num_physical_cores(); +} + +void sd_video_upscale_params_init(sd_video_upscale_params_t* params) { + if (!params) + return; + *params = {}; + params->fps = 24; + params->seed = 42; + params->window_size = 16; +} + +video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params) { + if (!params) + return nullptr; + try { + auto ctx = std::make_unique(); + if (!ctx->upscaler.init(*params)) + return nullptr; + return ctx.release(); + } catch (const std::exception& error) { + LOG_ERROR("video upscaler initialization failed: %s", error.what()); + return nullptr; + } +} + +bool upscale_video(video_upscaler_ctx_t* ctx, const sd_video_upscale_params_t* params, sd_video_generation_result_t** result_out) { + if (!result_out) + return false; + *result_out = nullptr; + if (!ctx || !params) + return false; + try { + return ctx->upscaler.upscale(*params, result_out); + } catch (const std::exception& error) { + LOG_ERROR("video upscale failed: %s", error.what()); + return false; + } +} + +void free_video_upscaler_ctx(video_upscaler_ctx_t* ctx) { + delete ctx; +} diff --git a/src/pipeline/video_upscale.h b/src/pipeline/video_upscale.h new file mode 100644 index 0000000000..62d0af1039 --- /dev/null +++ b/src/pipeline/video_upscale.h @@ -0,0 +1,46 @@ +#ifndef __SD_PIPELINE_VIDEO_UPSCALE_H__ +#define __SD_PIPELINE_VIDEO_UPSCALE_H__ + +#include + +#include "core/ggml_extend_backend.h" +#include "core/tensor.hpp" +#include "model_manager.h" +#include "stable-diffusion.h" + +namespace sd { + namespace pipeline { + + struct VideoUpscaleRuntime { + SDBackendManager backends; + std::shared_ptr manager; + int n_threads = 1; + }; + + class VideoUpscaleModel { + public: + virtual ~VideoUpscaleModel() = default; + virtual bool init(VideoUpscaleRuntime& runtime, const String2TensorStorage& weights, int upscale_factor) = 0; + virtual void end() = 0; + virtual bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t* result) = 0; + }; + + bool is_relay_vsr_model(const String2TensorStorage& weights); + std::unique_ptr create_relay_vsr_upscaler(); + + class VideoUpscaler { + VideoUpscaleRuntime runtime_; + int upscale_factor_ = 0; + std::unique_ptr executor_; + std::unique_ptr model_; + + public: + ~VideoUpscaler(); + bool init(const sd_video_upscaler_ctx_params_t& params); + bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t** result_out); + }; + + } // namespace pipeline +} // namespace sd + +#endif // __SD_PIPELINE_VIDEO_UPSCALE_H__ From 5f8c19f98ddb84e301bb8725b5260ec79e768e22 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 12:28:44 +0800 Subject: [PATCH 26/57] perf: reduce fused compute buffers (#2290) --- docs/ggml_extensions.md | 39 ++++++- ggml | 2 +- src/core/compute_workspace.cpp | 22 +++- src/core/cuda/convrot.cu | 3 +- src/core/cuda/cute_attention.cu | 86 +++++++++++---- src/core/cuda/i8_gated_residual.cu | 60 +++++++---- src/core/cuda/i8_gemm.cu | 93 ++++++++++++++--- src/core/cuda/i8_gemm.h | 3 + src/core/cuda/int8.cu | 161 ++++++++++++++++++++++++++--- src/core/cuda/rms_affine.cu | 34 ++++-- src/core/cuda/rope.cu | 5 +- 11 files changed, 423 insertions(+), 85 deletions(-) diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 77075989dd..580ef58984 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -89,8 +89,13 @@ 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. +consume nodes selected by an extension. Fused intermediates are excluded from +allocation through the backend's `skip_alloc` callback and +`ggml_gallocr_set_skipped_tensors`. Their metadata and graph edges remain available, +but execution must not access their data or buffers. Storage backing the final +output is retained, including an internal owner used by an in-place final write. +The skipped set is replaced for each allocation graph; changing which tensors are skipped +invalidates cached allocation assignments even when tensor shapes are unchanged. Automatic extension fusion is enabled in scheduler execution with one pipeline copy and without evaluation callbacks. Multiple pipeline copies and evaluation @@ -137,6 +142,17 @@ Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; caller must use the representation appropriate to their backend. Vulkan extensions use the command recording services described below. +CUDA INT8 Linear, RoPE, ConvRot, RMS affine and INT8 gated residual kernels omit output +scratch when allocation guarantees that a fresh owning output cannot overlap +live inputs. Existing outputs and views retain overlap protection. Internal +GEMM and quantization workspace is still reserved when required; kernels also +handle a null workspace when the entire step needs no scratch. + +Compute workspace measurement preserves the storage identities and offsets of +external tensors, including views. Separate allocations must not share a synthetic +buffer identity: that can invent overlap, overestimate scratch, and repeatedly +rebuild a workspace whose actual allocation is smaller than the estimate. + ## sd.cpp integration `src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas, @@ -202,6 +218,10 @@ ConvRot output fusion. These kernels use only the public extension API, borrowed CUDA stream, and scheduler-owned workspace; they do not use ggml's CUDA pool or private conversion helpers. The existing CUDA fusion and CuTe disable environment variables are honored. +Aligned contiguous FP16 K/V inputs are read directly without reserving conversion +buffers. Fresh F32 outputs are written by the attention epilogue, retaining the +original FP16 rounding before output scaling. Packed ConvRot and preallocated or +view outputs retain the intermediate half buffer. `src/core/cuda/i8_projection.cu` supplies the INT8 GEMM projection epilogues: `GEMM -> Reshape -> RMSNorm -> Mul -> RoPE`, optional scaling and FP16 copy, @@ -262,12 +282,21 @@ small-row policy, Stream-K scheduling, Ada column bands, and separate scaling order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row concatenation for custom GEMM nodes. Weight scales remain explicit inputs, including the gate/up scale permutation recorded on every SwiGLU launch. +Two independent INT8 GEMMs sharing a quantized input can also fuse with split +SwiGLU. A paired-column iterator reads the separate gate/up weight allocations, +and the GEMM epilogue writes only the activated output. Both intermediate +activations are omitted through the generic fusion allocation interface; weights +are neither copied nor concatenated. This path requires bias-free branches, +compatible layouts and alignments, and no external intermediate consumers. +Only interleaved weight scales and any Stream-K workspace remain temporary. Plans reserve bounded Stream-K workspace, cache occupancy on first execution, and initialize barriers on every launch, so CUDA Graph replay observes updated inputs and scales. External consumers retain ordinary execution; overlapping -outputs use workspace. 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. +outputs use workspace. For fresh outputs, the CUTLASS SwiGLU epilogue writes +an allocator-managed tensor that the subsequent ConvRot launch consumes. This +keeps the same GPU launches while allowing its storage to be reused by other +operations. The combined SwiGLU/ConvRot path retains its workspace intermediate +for other cases. Shapes or alignments unsuitable for CUTLASS use cuBLAS. Custom I8 nodes also retain cuBLAS when the reduction could overflow I32, to preserve their previous accumulation behavior. Fusion disable and force-cuBLAS settings are honored. No CUTLASS headers are needed for the cuBLAS path. diff --git a/ggml b/ggml index 135ce1e0e3..3fc478e619 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 135ce1e0e3f2b0693d3de93d4f6b5fa0a1947820 +Subproject commit 3fc478e6191aadf51ff7cf7bef15afe6ebeb14b1 diff --git a/src/core/compute_workspace.cpp b/src/core/compute_workspace.cpp index 4e531755b9..cf9d04c576 100644 --- a/src/core/compute_workspace.cpp +++ b/src/core/compute_workspace.cpp @@ -147,8 +147,8 @@ namespace sd { append(type); } } else { - const auto alignment_offset = reinterpret_cast(tensor->data) % ggml_backend_get_alignment(external); - append(alignment_offset); + append(tensor->buffer); + append(tensor->data); } const int32_t uses = ggml_graph_use_count(graph, tensor); append(uses); @@ -164,7 +164,7 @@ namespace sd { return {}; } std::unordered_map copies; - std::map external_buffers; + std::map, ggml_backend_buffer_t> external_buffers; for (auto tensor : tensors) { auto copy = ggml_dup_tensor(context, tensor); *copy = *tensor; @@ -179,7 +179,9 @@ namespace sd { } auto external = external_backend(source); if (external != nullptr && source->view_src == nullptr) { - auto& buffer = external_buffers[external]; + // Preserve distinct external allocations so overlap checks do not invent aliases. + const void* storage = source->buffer ? static_cast(source->buffer) : source; + auto& buffer = external_buffers[{external, storage}]; if (buffer == nullptr) { buffer = ggml_backend_alloc_buffer(external, 0); GGML_ASSERT(buffer != nullptr); @@ -191,6 +193,18 @@ namespace sd { copy->extra = nullptr; } } + for (const auto& entry : copies) { + const auto* source = entry.first; + if (source->view_src && source->buffer) { + const auto* owner = source->view_src; + while (owner->view_src) { + owner = owner->view_src; + } + if (external_backend(owner) != nullptr) { + entry.second->buffer = copies.at(owner)->buffer; + } + } + } auto copy_graph = ggml_new_graph_custom(context, graph_size, false); auto map_tensor = [](const ggml_tensor* tensor, void* user_data) -> ggml_tensor* { const auto& mapped = *static_cast*>(user_data); diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu index cf7896114d..26920ae934 100644 --- a/src/core/cuda/convrot.cu +++ b/src/core/cuda/convrot.cu @@ -204,7 +204,8 @@ namespace SD_CUDA_NAMESPACE { 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]); + const auto* dst = nodes[count - 1]; + *workspace = !dst->data && !dst->buffer && !dst->view_src ? 0 : ggml_nbytes(dst); return GGML_STATUS_SUCCESS; } static void destroy(void* state, void*) { diff --git a/src/core/cuda/cute_attention.cu b/src/core/cuda/cute_attention.cu index 1eaefa4b01..32bc891023 100644 --- a/src/core/cuda/cute_attention.cu +++ b/src/core/cuda/cute_attention.cu @@ -237,6 +237,9 @@ namespace sd_fattn_cute { const cutlass::half_t *q, *k, *v; cutlass::half_t* o; int nq, nk, heads, kv_heads, dim; + float* output_f32 = nullptr; + float output_scale = 1.0f, output_bias = 0.0f; + bool scaled_output = false; }; template @@ -333,18 +336,34 @@ namespace sd_fattn_cute { 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)); + auto output_regs = make_tensor(shape(from_sO)); 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); + const int64_t output_offset = int64_t(batch) * p.nq * p.heads * p.dim + int64_t(head) * p.dim; + if (p.output_f32) { + auto float_copy = make_tiled_copy(Copy_Atom, float>{}, + typename T::GlobalLayout{}, Layout>{}); + auto mO = make_tensor(make_gmem_ptr(p.output_f32 + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{})); + auto gO = local_tile(mO, Shape, Int>{}, make_coord(m_block, 0)); + auto to_gO = float_copy.get_thread_slice(tid).partition_D(gO); + auto values = make_tensor(shape(output_regs)); +#pragma unroll + for (int i = 0; i < size(values); ++i) { + // Preserve the half-rounded attention result before applying the F32 epilogue. + const float value = float(output_regs(i)); + values(i) = p.scaled_output ? __fmaf_rn(value, p.output_scale, p.output_bias) : value; + } + copy_rows(float_copy, values, to_gO, output_id, p.nq - m_block * BM, p.dim); + } else { + auto mO = make_tensor(make_gmem_ptr(p.o + output_offset), make_shape(p.nq, p.dim), make_stride(p.heads * p.dim, _1{})); + auto gO = local_tile(mO, Shape, Int>{}, make_coord(m_block, 0)); + auto to_gO = global_threadO.partition_D(gO); + copy_rows(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim); + } } } @@ -625,17 +644,34 @@ static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, vo struct AttentionPlan { int index; size_t offsets[4]; + bool direct_output; }; +static bool attention_direct_input(const ggml_tensor* src) { + size_t offset = 0; + const auto* owner = src; + while (!owner->data && owner->view_src) { + offset += owner->view_offs; + owner = owner->view_src; + } + return src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && (uintptr_t(owner->data) + offset) % 16 == 0; +} + static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) { auto* plan = new AttentionPlan{}; *state = plan; plan->index = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3 : 0; GGML_ASSERT(plan->index < count); - const auto* attn = nodes[plan->index]; - size_t size = 0; + const auto* attn = nodes[plan->index]; + const auto* dst = nodes[count - 1]; + plan->direct_output = dst->type == GGML_TYPE_F32 && !dst->data && !dst->buffer && !dst->view_src; + size_t size = 0; for (int i = 0; i < 4; ++i) { + if ((i > 0 && i < 3 && attention_direct_input(attn->src[i])) || (i == 3 && plan->direct_output)) { + plan->offsets[i] = SIZE_MAX; + continue; + } plan->offsets[i] = size; const size_t bytes = size_t(ggml_nelements(i < 3 ? attn->src[i] : attn)) * sizeof(half); size += (bytes + 255) / 256 * 256; @@ -681,7 +717,7 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution 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]); + work[i] = plan.offsets[i] == SIZE_MAX ? nullptr : 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]; @@ -695,6 +731,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution } 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 { + if (!work[i]) + return GGML_STATUS_FAILED; 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); @@ -702,12 +740,26 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution if (cudaGetLastError() != cudaSuccess) return GGML_STATUS_FAILED; } - const auto* q = attn->src[0]; - const auto* k = attn->src[1]; + const auto* q = attn->src[0]; + const auto* k = attn->src[1]; + auto* dst = nodes[count - 1]; + float scale = 1.0f, bias = 0.0f; + for (int i = plan.index + 1; i < count; ++i) { + if (nodes[i]->op == GGML_OP_SCALE) { + scale = param_f32(nodes[i], 0); + bias = param_f32(nodes[i], 1); + } + } sd_fattn_cute::Params p = {reinterpret_cast(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; + if (plan.direct_output) { + p.output_f32 = static_cast(dst->data); + p.output_scale = scale; + p.output_bias = bias; + p.scaled_output = dst->op == GGML_OP_SCALE; + } + cudaError_t error = cudaErrorInvalidValue; switch (padded_dim(p.dim)) { case 64: error = attention_launch<64>(context, p, q->ne[3]); @@ -733,14 +785,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution } 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]; + if (plan.direct_output) + return GGML_STATUS_SUCCESS; int convrot_group; if (sd_ggml_convrot_describe(dst, &convrot_group)) { const int64_t width = dst->src[0]->ne[0], rows = ggml_nrows(dst->src[0]), padded_rows = (rows + 3) / 4 * 4; diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu index 21b6134904..9e43d5c822 100644 --- a/src/core/cuda/i8_gated_residual.cu +++ b/src/core/cuda/i8_gated_residual.cu @@ -229,6 +229,7 @@ namespace sd_i8_gated { size_t kernel_workspace = 0; size_t output_offset = 0; int occupancy = 0; + bool output_scratch = true; ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr; }; template @@ -397,24 +398,50 @@ namespace sd_i8_gated { return configure(describe(nodes, count), plan) ? count : 0; #endif } + static uintptr_t address(const ggml_tensor* tensor) { + size_t offset = 0; + while (tensor && !tensor->data && tensor->view_src) { + offset += tensor->view_offs; + tensor = tensor->view_src; + } + return tensor && tensor->data ? reinterpret_cast(tensor->data) + offset : 0; + } + static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) { + if (!b) + return false; + const auto x = address(a), y = address(b); + if (!x || !y) + return false; + return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); + } + static bool needs_scratch(const ggml_tensor* dst, const Operation& op) { + if (address(dst) % 16 != 0) + return true; + for (const auto* input : {op.mat.weight, op.mat.input, op.mat.input_scale, op.mat.weight_scale, op.gate, op.residual, op.prefix}) { + if (input == op.residual && !op.prefix && address(input) == address(dst)) + continue; + if (overlaps(dst, input)) + return true; + } + return false; + } static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { auto* plan = new Plan; - *state = plan; - plan->device = device_info(ggml_backend_get_device(backend)); - if (!configure(describe(nodes, count), *plan)) + *state = plan; + plan->device = device_info(ggml_backend_get_device(backend)); + const auto op = describe(nodes, count); + if (!configure(op, *plan)) return GGML_STATUS_FAILED; - *workspace = plan->output_offset + ggml_nbytes(nodes[count - 1]); + // Fusion dependencies keep external inputs live until this output is allocated. + const auto* dst = nodes[count - 1]; + const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src; + plan->output_scratch = !fresh_output && !(dst->data && !needs_scratch(dst, op)); + *workspace = plan->output_offset + (plan->output_scratch ? ggml_nbytes(dst) : 0); return GGML_STATUS_SUCCESS; } static void destroy(void* state, void*) { delete static_cast(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); @@ -424,14 +451,11 @@ namespace sd_i8_gated { 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 stream = static_cast(context.stream); + auto* scratch = execution->workspace ? static_cast(ggml_backend_buffer_get_base(execution->workspace)) : nullptr; + const bool alias = needs_scratch(dst, op); + if (alias && !plan.output_scratch) + return GGML_STATUS_FAILED; auto* output = alias ? scratch + plan.output_offset : static_cast(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); diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu index 9d4d92e1ca..c62b95c776 100644 --- a/src/core/cuda/i8_gemm.cu +++ b/src/core/cuda/i8_gemm.cu @@ -105,6 +105,39 @@ namespace sd_i8_cutlass { 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 SplitWeightIterator : Base { + using Layout = typename Base::Layout; + using Pointer = typename Base::Pointer; + using TensorCoord = typename Base::TensorCoord; + using AccessType = typename Base::AccessType; + struct Params : Base::Params { + const int8_t* up = nullptr; + int64_t k = 0; + Params() = default; + CUTLASS_HOST_DEVICE Params(const Layout& layout) + : Base::Params(layout), k(layout.stride(0)) {} + }; + uintptr_t gate, up; + size_t gate_bytes; + CUTLASS_DEVICE SplitWeightIterator(const Params& params, Pointer pointer, TensorCoord extent, int thread, TensorCoord offset, const int* indices = nullptr) + : Base(params, pointer, extent, thread, offset, indices), + gate(reinterpret_cast(pointer)), + up(reinterpret_cast(params.up)), + gate_bytes(size_t(params.k) * (extent.column() / 2)) {} + + CUTLASS_DEVICE AccessType* get() const { + auto* address = Base::get(); + if (!address) + return nullptr; + // Keep the packed SwiGLU iterator's column pairing, but redirect its + // virtual up half to a separate allocation. Params.k stays full-sized + // even when Stream-K restricts this iterator's K extent. + const uintptr_t offset = reinterpret_cast(address) - gate; + return offset < gate_bytes ? address : reinterpret_cast(up + offset - gate_bytes); + } + }; + template struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore> { using Base = cutlass::epilogue::threadblock::VisitorAuxStore>; @@ -137,14 +170,15 @@ namespace sd_i8_cutlass { } }; - static __global__ void pair_scales(const float* src, float* dst, int n) { + static __global__ void pair_scales(const float* src, const float* up, float* dst, int n) { const int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) - dst[i] = src[i / 2 + (i % 2) * (n / 2)]; + dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)]; } template struct DenseTypes { + static constexpr bool kSplit = false; using Weight = std::conditional_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; @@ -182,8 +216,9 @@ namespace sd_i8_cutlass { std::conditional_t>::GemmKernel; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; - template + template struct SwigluTypes { + static constexpr bool kSplit = Split; using Weight = int8_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; @@ -198,9 +233,23 @@ namespace sd_i8_cutlass { 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 BaseMma = typename GemmBase::Mma; + using SplitMma = cutlass::gemm::threadblock::MmaMultistage< + typename BaseMma::Shape, + typename BaseMma::IteratorA, + typename BaseMma::SmemIteratorA, + BaseMma::kCacheOpA, + SplitWeightIterator, + typename BaseMma::SmemIteratorB, + BaseMma::kCacheOpB, + typename BaseMma::ElementC, + typename BaseMma::LayoutC, + typename BaseMma::Policy, + BaseMma::Detail::kStages>; + using Mma = std::conditional_t; using Kernel = std::conditional_t, - cutlass::gemm::kernel::GemmWithEpilogueVisitor>; + cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk, + cutlass::gemm::kernel::GemmWithEpilogueVisitor>; using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; struct Plan { @@ -208,9 +257,10 @@ namespace sd_i8_cutlass { int64_t rows, n, k; bool swiglu; bool packed_w4; + bool split; size_t kernel_bytes = 0, scratch_offset = 0; int occupancy = 0; - ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t) = nullptr; + ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr; }; template 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) { @@ -219,10 +269,10 @@ namespace sd_i8_cutlass { return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast(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) { + static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales) { using Kernel = typename Types::Kernel; if ((reinterpret_cast(x) | reinterpret_cast(w) | reinterpret_cast(xs) | - reinterpret_cast(ws) | reinterpret_cast(out)) % + reinterpret_cast(ws) | reinterpret_cast(out) | reinterpret_cast(up) | reinterpret_cast(up_scales)) % 16) return GGML_STATUS_FAILED; auto* paired = static_cast(scratch); @@ -239,14 +289,19 @@ namespace sd_i8_cutlass { return GGML_STATUS_FAILED; } typename Kernel::Params params(args, p.device.sms, p.occupancy); + if constexpr (Types::kSplit) { + if (!up || !up_scales) + return GGML_STATUS_FAILED; + params.params_B.up = up; + } if (params.get_workspace_size() > p.kernel_bytes) return GGML_STATUS_FAILED; if (p.swiglu) { - pair_scales<<>>(ws, paired, int(p.n)); + pair_scales<<>>(ws, up_scales, 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) + if (params.init_workspace(scratch ? static_cast(scratch) + p.scratch_offset : nullptr, stream) != cutlass::Status::kSuccess) return GGML_STATUS_FAILED; cutlass::arch::synclog_setup(); cutlass::Kernel2<<>>(params); @@ -277,6 +332,8 @@ namespace sd_i8_cutlass { static bool configure_tile(Plan& p) { if (p.packed_w4) return configure_kernel>(p); + if (p.split) + return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); if (p.swiglu) return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); return p.rows >= 8192 ? configure_kernel>(p) : configure_kernel>(p); @@ -305,7 +362,7 @@ namespace sd_i8_cutlass { } } 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, bool packed_w4) { +static SDI8CudaGemmPlan* prepare_gemm(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4, bool split) { #ifdef SD_CUDA_FORCE_CUBLAS return nullptr; #else @@ -320,6 +377,7 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row p->k = k; p->swiglu = swiglu; p->packed_w4 = packed_w4; + p->split = split; if (!sd_i8_cutlass::configure(*p)) { delete p; return nullptr; @@ -328,8 +386,19 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t row return p; #endif } +SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4) { + return prepare_gemm(device, rows, n, k, swiglu, workspace, packed_w4, false); +} +SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace) { + if (n > INT_MAX / 2 || k > INT_MAX / 16384) + return nullptr; + return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true); +} ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) { - return plan->launch(*plan, x, w, xs, ws, out, scratch, stream); + return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr); +} +ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) { + return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales); } void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan) { delete plan; diff --git a/src/core/cuda/i8_gemm.h b/src/core/cuda/i8_gemm.h index 5e4ac9f8c2..dc0d15d6ea 100644 --- a/src/core/cuda/i8_gemm.h +++ b/src/core/cuda/i8_gemm.h @@ -5,5 +5,8 @@ struct SDI8CudaGemmPlan; SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false); ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream); +// n is the width of each branch; gate/up weights and scales remain separate. +SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace); +ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream); void sd_cuda_i8_gemm_free(SDI8CudaGemmPlan* plan); #endif diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 60f36bb1e9..93d8f7f0d0 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -445,6 +445,13 @@ namespace SD_CUDA_NAMESPACE { static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }(); return disabled; } +#ifdef SD_CUDA_I8_CUTLASS + static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) { + return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 && + aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) && + (op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))); + } +#endif static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { auto* g = const_cast(graph); if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g)) @@ -466,8 +473,21 @@ namespace SD_CUDA_NAMESPACE { const auto* quant = ggml_graph_node(g, first + 2); int group; if (quant->src[0] == next && sd_ggml_convrot_describe(quant, &group) && group == 256 && convrot_supports(device, quant, nullptr) && - size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) + size_t(next->ne[0]) * sizeof(float) + 128 <= cuda_device_info(device).shared_memory) { +#ifdef SD_CUDA_I8_CUTLASS + if (!next->data && !next->buffer && !next->view_src && cutlass_compatible(op, next)) { + size_t workspace = 0; + auto* gemm = sd_cuda_i8_gemm_prepare(device, ggml_nrows(mat), mat->ne[0], op.weight->ne[0], true, &workspace); + const bool fused_epilogue = gemm != nullptr; + sd_cuda_i8_gemm_free(gemm); + // ConvRot is a separate launch either way. Keep its input in the + // tensor arena so other operations can reuse this large allocation. + if (fused_epilogue) + return 2; + } +#endif return 3; + } } return 2; } @@ -511,9 +531,7 @@ namespace SD_CUDA_NAMESPACE { #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))) && + if (cutlass_compatible(op, dst) && (p->kind != CONCAT || ggml_nbytes(dst->src[0]) % 16 == 0)) { if (p->kind == SWIGLU || p->kind == SWIGLU_QUANTIZE) { p->gemm = sd_cuda_i8_gemm_prepare(device, p->rows, p->n, p->k, true, &gemm_bytes); @@ -530,8 +548,12 @@ namespace SD_CUDA_NAMESPACE { const bool large_rotation = (!mat || (op.input->type == GGML_TYPE_F32 && op.group_size == 256)) && size_t(p->k) * sizeof(float) + 128 > p->device.shared_memory; p->partial = reserve(total, p->group == 64 ? size_t(p->rows) * (p->k / 64) * sizeof(float) : large_rotation ? size_t(p->rows) * (p->k / 256) * sizeof(float) : 0); - p->output = reserve(total, p->gemm && p->kind == SWIGLU_QUANTIZE ? 0 : ggml_nbytes(dst)); - *workspace = total; + // Custom/concat outputs and packed GLU cannot reuse their input storage. + const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src && + (dst->op == GGML_OP_CUSTOM || dst->op == GGML_OP_CONCAT || + (dst->op == GGML_OP_GLU && !dst->src[1])); + p->output = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst)); + *workspace = total; return GGML_STATUS_SUCCESS; } static void destroy(void* state, void*) { @@ -548,6 +570,99 @@ namespace SD_CUDA_NAMESPACE { 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); } +#ifdef SD_CUDA_I8_CUTLASS + static bool split_describe(ggml_backend_dev_t device, ggml_tensor* const* nodes, SDI8Operation& gate, SDI8Operation& up) { + const auto* glu = nodes[2]; + if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || glu->op_params[1] != 0 || + !glu->src[0] || !glu->src[1] || glu->src[0] == glu->src[1] || glu->view_src || + glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) || !aligned(glu, 16) || + !((nodes[0] == glu->src[0] && nodes[1] == glu->src[1]) || (nodes[1] == glu->src[0] && nodes[0] == glu->src[1]))) + return false; + for (const auto* mat : {glu->src[0], glu->src[1]}) + if (mat->view_src || !supports(device, mat, nullptr) || !ggml_are_same_shape(mat, glu)) + return false; + if (!sd_ggml_i8_describe(glu->src[0], &gate) || !sd_ggml_i8_describe(glu->src[1], &up) || + gate.input != up.input || gate.input_scale != up.input_scale || gate.packed != up.packed || + gate.input->type != GGML_TYPE_I8 || gate.bias || up.bias || !gate.weight_scale || !up.weight_scale || + !ggml_are_same_shape(gate.weight, up.weight)) + return false; + for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale}) + if (!aligned(t, 16)) + return false; + return true; + } + static int split_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) { + auto* g = const_cast(graph); + if (fusion_disabled() || first + 2 >= ggml_graph_n_nodes(g)) + return 0; + ggml_tensor* nodes[] = {ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)}; + SDI8Operation gate, up; + if (!split_describe(device, nodes, gate, up)) + return 0; + size_t workspace = 0; + auto* gemm = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(nodes[2]), nodes[2]->ne[0], gate.weight->ne[0], &workspace); + const bool supported = gemm != nullptr; + sd_cuda_i8_gemm_free(gemm); + return supported ? 3 : 0; + } + struct SplitPlan { + SDI8CudaGemmPlan* gemm = nullptr; + size_t output = SIZE_MAX; + int device; + }; + static ggml_status split_prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { + auto* p = new SplitPlan; + *state = p; + auto device = ggml_backend_get_device(backend); + p->device = cuda_device_info(device).device; + SDI8Operation gate, up; + if (count != 3 || !split_describe(device, nodes, gate, up)) + return GGML_STATUS_FAILED; + auto* dst = nodes[2]; + size_t bytes = 0; + p->gemm = sd_cuda_i8_gemm_prepare_split(device, ggml_nrows(dst), dst->ne[0], gate.weight->ne[0], &bytes); + if (!p->gemm) + return GGML_STATUS_FAILED; + size_t total = (bytes + 255) / 256 * 256; + if (dst->data || dst->buffer || dst->view_src) + p->output = reserve(total, ggml_nbytes(dst)); + *workspace = total; + return GGML_STATUS_SUCCESS; + } + static void split_destroy(void* state, void*) { + auto* p = static_cast(state); + if (p) + sd_cuda_i8_gemm_free(p->gemm); + delete p; + } + static ggml_status split_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void* state, void*) { + auto& p = *static_cast(state); + auto* dst = nodes[2]; + SDI8Operation gate, up; + if (!sd_ggml_i8_describe(dst->src[0], &gate) || !sd_ggml_i8_describe(dst->src[1], &up)) + return GGML_STATUS_FAILED; + ggml_ext_gpu_context context; + if (!cuda_context(execution->backend, &context) || context.device != p.device) + return GGML_STATUS_FAILED; + auto stream = static_cast(context.stream); + auto* scratch = static_cast(ggml_backend_buffer_get_base(execution->workspace)); + bool alias = false; + for (const auto* t : {gate.input, gate.input_scale, gate.weight, up.weight, gate.weight_scale, up.weight_scale}) + alias |= overlaps(dst, t); + if (alias && p.output == SIZE_MAX) + return GGML_STATUS_FAILED; + auto* output = alias ? reinterpret_cast(scratch + p.output) : static_cast(dst->data); + const auto* input = static_cast(gate.input->data); + const int64_t padded = (ggml_nrows(dst) + 3) / 4 * 4; + const auto* scales = gate.packed ? reinterpret_cast(input + gate.weight->ne[0] * padded) : static_cast(gate.input_scale->data); + auto status = sd_cuda_i8_gemm_execute_split(p.gemm, input, static_cast(gate.weight->data), static_cast(up.weight->data), + scales, static_cast(gate.weight_scale->data), static_cast(up.weight_scale->data), + output, scratch, stream); + if (status != GGML_STATUS_SUCCESS || (alias && cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream) != cuda_success)) + return GGML_STATUS_FAILED; + return GGML_STATUS_SUCCESS; + } +#endif static CUDAError convrot_launch(const Plan& p, CUDAStream stream, const float* src, int8_t* dst, float* scales, float* partial) { if (p.padded > p.rows) { auto error = cuda_memset_async(dst + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); @@ -642,9 +757,10 @@ namespace SD_CUDA_NAMESPACE { 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]; + auto stream = static_cast(context.stream); + auto* scratch = execution->workspace ? static_cast(ggml_backend_buffer_get_base(execution->workspace)) : nullptr; + auto scratch_at = [scratch](size_t offset) { return scratch ? scratch + offset : nullptr; }; + auto* dst = nodes[count - 1]; SDI8Operation op; const bool mat = sd_ggml_i8_describe(nodes[0], &op); bool alias = false; @@ -659,20 +775,22 @@ namespace SD_CUDA_NAMESPACE { // 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; + if (alias && p.output == SIZE_MAX) + return GGML_STATUS_FAILED; + void* output = alias ? scratch_at(p.output) : dst->data; CUDAError error = cuda_success; if (!mat) { auto* quant = static_cast(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)); + error = convrot_launch(p, stream, static_cast(dst->src[0]->data), quant, scales, reinterpret_cast(scratch_at(p.partial))); } else { - auto* quant = reinterpret_cast(scratch + p.quant); + auto* quant = reinterpret_cast(scratch_at(p.quant)); const float* scales; if (op.input->type == GGML_TYPE_F32) { - auto* temporary_scales = reinterpret_cast(scratch + p.scales); + auto* temporary_scales = reinterpret_cast(scratch_at(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)); + error = convrot_launch(p, stream, static_cast(op.input->data), quant, temporary_scales, reinterpret_cast(scratch_at(p.partial))); else { if (p.padded > p.rows) error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream); @@ -696,12 +814,12 @@ namespace SD_CUDA_NAMESPACE { 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); + auto* accum = reinterpret_cast(scratch_at(p.accum)); const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu); #ifdef SD_CUDA_I8_CUTLASS if (p.gemm) { auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast(op.weight->data), scales, ws, - intermediate ? reinterpret_cast(accum) : gemm_output, scratch + p.gemm_offset, stream); + intermediate ? reinterpret_cast(accum) : gemm_output, scratch_at(p.gemm_offset), stream); if (status != GGML_STATUS_SUCCESS) return status; } else @@ -794,6 +912,17 @@ namespace SD_CUDA_NAMESPACE { desc.custom_op = sd_ggml_convrot_op(); desc.supports = sd_i8_cuda::convrot_supports; GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); +#ifdef SD_CUDA_I8_CUTLASS + desc.name = "sd.linear_i8.swiglu_split.cuda.v1"; + desc.custom_op = sd_ggml_i8_op(); + desc.priority = 260; + desc.supports = nullptr; + desc.match = sd_i8_cuda::split_match; + desc.prepare = sd_i8_cuda::split_prepare; + desc.execute = sd_i8_cuda::split_execute; + desc.destroy = sd_i8_cuda::split_destroy; + GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE)); +#endif } } diff --git a/src/core/cuda/rms_affine.cu b/src/core/cuda/rms_affine.cu index 9f4192b259..02196ffa8d 100644 --- a/src/core/cuda/rms_affine.cu +++ b/src/core/cuda/rms_affine.cu @@ -169,15 +169,35 @@ namespace SD_CUDA_NAMESPACE { } } + static uintptr_t address(const ggml_tensor* tensor) { + size_t offset = 0; + while (tensor && !tensor->data && tensor->view_src) { + offset += tensor->view_offs; + tensor = tensor->view_src; + } + return tensor && tensor->data ? reinterpret_cast(tensor->data) + offset : 0; + } 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); + const auto x = address(a), y = address(b); + if (!x || !y) + return false; return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b); } + static bool needs_scratch(const Operation& op) { + const auto* dst = op.output; + return (overlaps(dst, op.input) && address(dst) != address(op.input)) || + (overlaps(dst, op.residual) && address(dst) != address(op.residual)) || + overlaps(dst, op.weight) || overlaps(dst, op.gate) || overlaps(dst, op.shift); + } static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) { - *workspace = ggml_nbytes(nodes[count - 1]); + const auto op = describe(nodes, count); + const auto* dst = op.output; + const bool fresh = !dst->data && !dst->buffer && !dst->view_src; + // Fresh outputs may reuse the residual elementwise; external inputs stay live. + *workspace = fresh || (dst->data && !needs_scratch(op)) ? 0 : ggml_nbytes(dst); return GGML_STATUS_SUCCESS; } static void destroy(void*, void*) {} @@ -206,11 +226,11 @@ namespace SD_CUDA_NAMESPACE { 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); + const auto op = describe(nodes, count); + auto* dst = op.output; + const bool alias = needs_scratch(op); + if (alias && !execution->workspace) + return GGML_STATUS_FAILED; auto* output = static_cast(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data); const auto stream = static_cast(context.stream); CUDAError error = cuda_invalid_value; diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index 99cc824f1f..f3ab216e67 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -277,7 +277,10 @@ namespace SD_CUDA_NAMESPACE { } 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]); + const auto* dst = nodes[count - 1]; + // Owning casts are fresh outputs; copies into existing storage retain view_src. + const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src; + *workspace = fresh_output ? 0 : ggml_nbytes(dst); return GGML_STATUS_SUCCESS; } From 7bcd4ce8ce69378115c176d46d4f30aed80c005a Mon Sep 17 00:00:00 2001 From: Zihao Date: Sun, 8 Nov 2026 18:23:11 +1100 Subject: [PATCH 27/57] fix: clean LoRA multiplier and prompt spacing in image metadata (#2283) --- examples/common/generation_metadata.cpp | 13 +++++++++++-- src/generation_info.cpp | 25 +++++++++++++++++++++++-- 2 files changed, 34 insertions(+), 4 deletions(-) diff --git a/examples/common/generation_metadata.cpp b/examples/common/generation_metadata.cpp index 989e914462..05205fffdb 100644 --- a/examples/common/generation_metadata.cpp +++ b/examples/common/generation_metadata.cpp @@ -1,5 +1,6 @@ #include "generation_metadata.h" +#include #include #include @@ -70,6 +71,7 @@ std::string generation_info_to_webui(const std::string& text) { const auto& sampling = info.at("sample_params"); std::string result = info.value("prompt", ""); if (info.contains("loras")) { + bool prompt_trimmed = false; for (const auto& lora : info.at("loras")) { // WebUI tags cannot restrict a LoRA to the high-noise model. if (lora.value("is_high_noise", false)) @@ -83,8 +85,15 @@ std::string generation_info_to_webui(const std::string& text) { name.erase(extension); if (name.empty() || name.find_first_of(":<>\r\n") != std::string::npos) continue; - if (!result.empty()) - result += " "; + if (!result.empty()) { + if (!prompt_trimmed) { + while (!result.empty() && std::isspace(static_cast(result.back()))) + result.pop_back(); + prompt_trimmed = true; + } + if (!result.empty()) + result += " "; + } result += ""; } } diff --git a/src/generation_info.cpp b/src/generation_info.cpp index cf40bc2bad..0d4829139d 100644 --- a/src/generation_info.cpp +++ b/src/generation_info.cpp @@ -1,6 +1,7 @@ #include "generation_info.h" #include +#include #include #include @@ -14,6 +15,21 @@ namespace sd { return value ? value : ""; } + static double shortest_f32_as_double(float value) { + if (!std::isfinite(value)) { + return static_cast(value); + } + char buf[64]; + buf[0] = '\0'; + for (int prec = 1; prec <= 9; ++prec) { + std::snprintf(buf, sizeof(buf), "%.*g", prec, static_cast(value)); + if (std::strtof(buf, nullptr) == value) { + break; + } + } + return std::strtod(buf, nullptr); + } + 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}}}}; @@ -117,8 +133,13 @@ namespace sd { } 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}}); + for (uint32_t i = 0; i < p.lora_count; ++i) { + // nlohmann::json promotes float to double, so dump the shortest + // decimal that round-trips the float32 (e.g. 0.85f stays 0.85 + // instead of 0.8500000238418579) to keep replay exact. + const double multiplier = shortest_f32_as_double(p.loras[i].multiplier); + info["loras"].push_back({{"path", text(p.loras[i].path)}, {"multiplier", multiplier}, {"is_high_noise", p.loras[i].is_high_noise}}); + } } return info; } From de00bafc3d473069583752c80a0d56140477ede2 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 15:50:25 +0800 Subject: [PATCH 28/57] feat: add configurable attention chunk budget for MiniMax-H3 (#2292) --- docs/minimax_h3.md | 12 ++++++ src/core/cuda/i8_gated_residual.cu | 2 +- src/core/cuda/i8_gemm.cu | 62 +++++++++++++-------------- src/core/cuda/int8.cu | 4 +- src/core/cuda/rope.cu | 2 +- src/model/diffusion/minimax_h3.hpp | 61 +++++++++++++++++++++++--- src/model/integrations/minimax_h3.cpp | 3 +- 7 files changed, 103 insertions(+), 43 deletions(-) diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index db554ddb03..bc1f19c770 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -155,5 +155,17 @@ Notes for this mode: - Width and height are aligned upward to a multiple of 32. - Frame count is aligned upward to the `17k + 5` grid, with a minimum of 5. - MiniMax-H3 runs at 24 fps; another requested value is overridden. +- Attention head grouping is controlled by + `--model-args minimax_h3_attn_chunk_budget=2`, in GiB (default: `2`). + The budget estimates the size of a single FP32 Q/K tensor per group: + `head_dim * tokens * batch * 4 bytes * heads_per_group`. The group size adapts + to the sequence length, with a minimum of one head; inputs that fit use all + heads together. Set `0` to disable grouping, or a smaller budget such as + `0.25` to reduce peak activation memory for long sequences at the cost of + extra launches and concatenation. Each head still attends to the full sequence. + This is a grouping target, not a cap on total attention memory or VRAM usage; + other intermediates, backend workspace, and weights require additional memory. + `--max-vram` and `--offload-to-cpu` control weight residency, but the active + block's working memory must also fit within the VRAM budget. - The default video flow shift is 12. The audio stream is mapped internally to its shift of 3, so the regular samplers can operate on the packed AV latent. diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu index 9e43d5c822..71477d7432 100644 --- a/src/core/cuda/i8_gated_residual.cu +++ b/src/core/cuda/i8_gated_residual.cu @@ -426,7 +426,7 @@ namespace sd_i8_gated { return false; } static ggml_status prepare(ggml_backend_t backend, ggml_tensor* const* nodes, int count, void** state, size_t* workspace, void*) { - auto* plan = new Plan; + auto* plan = new Plan; *state = plan; plan->device = device_info(ggml_backend_get_device(backend)); const auto op = describe(nodes, count); diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu index c62b95c776..9d06c91581 100644 --- a/src/core/cuda/i8_gemm.cu +++ b/src/core/cuda/i8_gemm.cu @@ -179,19 +179,19 @@ namespace sd_i8_cutlass { template struct DenseTypes { static constexpr bool kSplit = false; - using Weight = std::conditional_t; - 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< + using Weight = std::conditional_t; + 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, @@ -219,20 +219,20 @@ namespace sd_i8_cutlass { template struct SwigluTypes { static constexpr bool kSplit = Split; - using Weight = int8_t; - using Tile = typename Config::Tile; - using Warp = typename Config::Warp; - using Swizzle = std::conditional_t>; - 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 Weight = int8_t; + 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 BaseMma = typename GemmBase::Mma; using SplitMma = cutlass::gemm::threadblock::MmaMultistage< typename BaseMma::Shape, @@ -246,11 +246,11 @@ namespace sd_i8_cutlass { typename BaseMma::LayoutC, typename BaseMma::Policy, BaseMma::Detail::kStages>; - using Mma = std::conditional_t; - using Kernel = std::conditional_t; + using Kernel = std::conditional_t, cutlass::gemm::kernel::GemmWithEpilogueVisitor>; - using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; }; struct Plan { DeviceInfo device; @@ -259,7 +259,7 @@ namespace sd_i8_cutlass { bool packed_w4; bool split; size_t kernel_bytes = 0, scratch_offset = 0; - int occupancy = 0; + int occupancy = 0; ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr; }; template diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 93d8f7f0d0..01a614731e 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -552,8 +552,8 @@ namespace SD_CUDA_NAMESPACE { const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src && (dst->op == GGML_OP_CUSTOM || dst->op == GGML_OP_CONCAT || (dst->op == GGML_OP_GLU && !dst->src[1])); - p->output = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst)); - *workspace = total; + p->output = fresh_output || (p->gemm && p->kind == SWIGLU_QUANTIZE) ? SIZE_MAX : reserve(total, ggml_nbytes(dst)); + *workspace = total; return GGML_STATUS_SUCCESS; } static void destroy(void* state, void*) { diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index f3ab216e67..2bad5df52a 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -277,7 +277,7 @@ namespace SD_CUDA_NAMESPACE { } static ggml_status rope_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) { - const auto* dst = nodes[count - 1]; + const auto* dst = nodes[count - 1]; // Owning casts are fresh outputs; copies into existing storage retain view_src. const bool fresh_output = !dst->data && !dst->buffer && !dst->view_src; *workspace = fresh_output ? 0 : ggml_nbytes(dst); diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 9a40643aab..97fd230e3e 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -19,9 +19,10 @@ namespace MiniMaxH3 { - constexpr int H3_GRAPH_SIZE = 131072; - constexpr float FRAME_RESCALE = 5.f / 3.f; - constexpr float VISUAL_COND_TIMESTEP = 0.999f; + constexpr int H3_GRAPH_SIZE = 131072; + constexpr float DEFAULT_ATTN_CHUNK_BUDGET_GIB = 2.f; + constexpr float FRAME_RESCALE = 5.f / 3.f; + constexpr float VISUAL_COND_TIMESTEP = 0.999f; struct Config { int64_t hidden_size = 5376; @@ -231,6 +232,7 @@ namespace MiniMaxH3 { struct Attention : public GGMLBlock { int64_t heads; int64_t head_dim; + float head_chunk_budget_gib = DEFAULT_ATTN_CHUNK_BUDGET_GIB; Attention(int64_t hidden_size, int64_t heads, @@ -315,6 +317,31 @@ namespace MiniMaxH3 { ggml_tensor* attend(GGMLRunnerContext* ctx, const std::vector& qkv, ggml_tensor* pe = nullptr) { + // Budget one F32 Q/K tensor per group, not the full attention working set. + const double head_bytes = double(ggml_row_size(GGML_TYPE_F32, head_dim)) * qkv[0]->ne[2] * qkv[0]->ne[3]; + const int64_t head_chunk = head_chunk_budget_gib > 0 + ? static_cast(std::clamp(double(head_chunk_budget_gib) * (1024ULL * 1024 * 1024) / head_bytes, + 1.0, double(heads))) + : heads; + if (head_chunk < heads) { + std::vector outputs; + for (int64_t start = 0; start < heads; start += head_chunk) { + std::vector part; + for (auto tensor : qkv) { + part.push_back(ggml_ext_slice(ctx->ggml_ctx, tensor, 1, start, std::min(start + head_chunk, heads), false)); + } + outputs.push_back(attend_heads(ctx, part, pe)); + // Keep each head group's temporaries out of the concatenation tree's live set. + ctx->expand_graph(outputs.back()); + } + return ggml_ext_vec_concat(ctx->ggml_ctx, outputs, 0); + } + return attend_heads(ctx, qkv, pe); + } + + ggml_tensor* attend_heads(GGMLRunnerContext* ctx, + const std::vector& qkv, + ggml_tensor* pe) { auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); auto q = qkv[0]; @@ -333,7 +360,7 @@ namespace MiniMaxH3 { q, k, v, - static_cast(heads), + static_cast(qkv[0]->ne[1]), nullptr, true, ctx->flash_attn_enabled, @@ -1035,10 +1062,30 @@ namespace MiniMaxH3 { MiniMaxH3Runner(ggml_backend_t backend, const String2TensorStorage& tensors, const std::string& prefix = "model.diffusion_model", - std::shared_ptr weight_manager = nullptr) + std::shared_ptr weight_manager = nullptr, + const char* model_args = nullptr) : DiffusionModelRunner(backend, prefix, weight_manager), config(Config::detect_from_weights(tensors, prefix)), model(config) { + float head_chunk_budget_gib = DEFAULT_ATTN_CHUNK_BUDGET_GIB; + for (const auto& [key, value] : parse_key_value_args(model_args, "model arg")) { + if (key == "minimax_h3_attn_chunk_budget") { + float parsed = 0.f; + if (parse_strict_float(value, parsed) && std::isfinite(parsed) && parsed >= 0.f) { + head_chunk_budget_gib = parsed; + } else { + LOG_WARN("ignoring invalid MiniMax-H3 model arg '%s=%s': expected a finite non-negative GiB budget", key.c_str(), value.c_str()); + } + } + } + std::vector blocks; + model.get_all_blocks(blocks); + for (auto block : blocks) { + if (auto attention = dynamic_cast(block)) { + attention->head_chunk_budget_gib = head_chunk_budget_gib; + } + } + LOG_INFO("MiniMax-H3 attention chunk budget: %.3g GiB (0 disables grouping)", head_chunk_budget_gib); model.init(params_ctx, tensors, prefix); } @@ -1178,7 +1225,8 @@ namespace MiniMaxH3 { timestep_features = make_input(timestep_feature_input_cache); } - auto runner_ctx = get_context(); + auto graph = new_graph_custom(H3_GRAPH_SIZE); + auto runner_ctx = get_context(graph); auto output = model.forward(&runner_ctx, video, audio, @@ -1201,7 +1249,6 @@ namespace MiniMaxH3 { output.second, 1.f + (audio_scale - 1.f) * sigma_a)); auto merged = merge_av_latents(compute_ctx, output.first, output.second); - auto graph = new_graph_custom(H3_GRAPH_SIZE); ggml_build_forward_expand(graph, merged); return graph; } diff --git a/src/model/integrations/minimax_h3.cpp b/src/model/integrations/minimax_h3.cpp index 35a1f6700b..7624a6cc13 100644 --- a/src/model/integrations/minimax_h3.cpp +++ b/src/model/integrations/minimax_h3.cpp @@ -18,7 +18,8 @@ static bool build_minimax_h3(const sd::model_builders::Context& ctx, const Token result.diffusion = std::make_shared(ctx.backends.runtime_backend(SDBackendModule::DIFFUSION), tensor_storage_map, "model.diffusion_model", - weight_manager); + weight_manager, + sd_ctx_params->model_args); return true; } From 0b3b98d7b83ec433925252e1fbc00118629a3c5f Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 16:24:53 +0800 Subject: [PATCH 29/57] fix: avoid CPU fallback based on backend allocation chunk size (#2294) --- src/core/ggml_runner.cpp | 7 +++++-- src/model/diffusion/minimax_h3.hpp | 23 ++++++++++++----------- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 2fa14b1508..770f46b219 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -212,6 +212,10 @@ ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const { } static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { + // Vulkan reports a preferred allocation chunk size, not a hard tensor limit. + if (sd_backend_is(backend, "Vulkan")) { + return true; + } auto buft = ggml_backend_get_default_buffer_type(backend); size_t alignment = ggml_backend_buft_get_alignment(buft); size_t max_size = ggml_backend_buft_get_max_size(buft); @@ -219,6 +223,7 @@ static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* n if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { return false; } + // Cross-backend copies must also fit the destination's allocation limit. for (auto src : node->src) { if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { return false; @@ -251,8 +256,6 @@ void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { continue; } - // Operation support can allow tensors larger than the backend's allocation limit. - // Check inputs too, since cross-backend copies must also fit a single allocation. if (!backend_can_allocate_op(current, node)) { auto cpu = workspace_.cpu_backend(); if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 97fd230e3e..5f388d84f4 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -174,14 +174,14 @@ namespace MiniMaxH3 { 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. + auto buft = ggml_backend_get_default_buffer_type(ctx->backend); + size_t alignment = ggml_backend_buft_get_alignment(buft); + size_t chunk_budget = ggml_backend_buft_get_max_size(buft); + chunk_budget -= chunk_budget % alignment; + // Use the backend's preferred allocation size as a conservative chunk budget. + // Matmul outputs are F32 even for quantized weights. 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))); + chunk_size = static_cast(std::min(chunk_size, std::max(1, chunk_budget / token_bytes))); } if (chunk_size == x->ne[1]) { return forward_chunk(x); @@ -269,10 +269,11 @@ namespace MiniMaxH3 { 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; + // This is a conservative chunk budget, not an operation support limit. + size_t chunk_budget = ggml_backend_buft_get_max_size(buft); + chunk_budget -= chunk_budget % alignment; size_t token_bytes = ggml_row_size(GGML_TYPE_F32, 3 * heads * head_dim) * inputs[0]->ne[2] * inputs[0]->ne[3]; - chunk_size = static_cast(std::min(chunk_size, std::max(1, max_size / token_bytes))); + chunk_size = static_cast(std::min(chunk_size, std::max(1, chunk_budget / token_bytes))); } if (chunk_size == sequence) { ggml_tensor* projected = nullptr; @@ -284,7 +285,7 @@ namespace MiniMaxH3 { return split_qkv(ctx, projected); } - // Split Q/K/V before concatenating tokens so the full packed QKV never exceeds the buffer limit. + // Split Q/K/V before concatenating tokens to avoid materializing the full packed QKV. std::vector chunks[3]; for (size_t i = 0; i < inputs.size(); ++i) { auto input = inputs[i]; From f61053c25d37e7dc91a8525febb5fd10a68a4787 Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 20:07:05 +0800 Subject: [PATCH 30/57] fix(metal): handle all batch elements in group normalization (#2296) --- ggml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml b/ggml index 3fc478e619..a70d9e12b3 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 3fc478e6191aadf51ff7cf7bef15afe6ebeb14b1 +Subproject commit a70d9e12b381f778bbb92505ecbe370469b43a52 From b43af40461cc64c89cac3b621019f8914f50dd9c Mon Sep 17 00:00:00 2001 From: DrMkdaddy <37319157+DrMkdaddy@users.noreply.github.com> Date: Sun, 8 Nov 2026 05:37:06 -0800 Subject: [PATCH 31/57] fix: use patch token count for the Flux scheduler (#2293) --- src/model/diffusion/flux.hpp | 4 ++++ src/model/diffusion/model.hpp | 4 ++++ src/pipeline/diffusion_engine.cpp | 7 ++++++- 3 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index 843e70b458..f8f8305dd4 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -1437,6 +1437,10 @@ namespace Flux { return "flux"; } + int get_patch_size() const override { + return config.patch_size; + } + void get_param_tensors(std::map& tensors, const std::string& prefix) override { flux.get_param_tensors(tensors, prefix); } diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index 4ff91c416c..f25ca7b519 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -237,6 +237,10 @@ struct DiffusionModelRunner : public GGMLRunner { virtual sd::Tensor compute(int n_threads, const DiffusionParams& diffusion_params) = 0; + virtual int get_patch_size() const { + return 1; + } + void get_param_tensors(std::map& tensors) { get_param_tensors(tensors, prefix); } diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index a8eed8e771..c820a8d963 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2930,7 +2930,12 @@ int StableDiffusionGGML::get_image_channels() const { int StableDiffusionGGML::get_image_seq_len(int h, int w) { int vae_scale_factor = get_vae_scale_factor(); - return (h / vae_scale_factor) * (w / vae_scale_factor); + int patch_size = diffusion_model != nullptr ? diffusion_model->get_patch_size() : 1; + int latent_h = h / vae_scale_factor; + int latent_w = w / vae_scale_factor; + int tokens_h = (latent_h + patch_size - 1) / patch_size; + int tokens_w = (latent_w + patch_size - 1) / patch_size; + return tokens_h * tokens_w; } sd::Tensor StableDiffusionGGML::generate_init_latent(int width, From 3d315c47d4ed5c60d8d686a0ba7e3ca09147141b Mon Sep 17 00:00:00 2001 From: leejet Date: Sun, 8 Nov 2026 22:28:35 +0800 Subject: [PATCH 32/57] feat: add Metal INT8 ConvRot support (#2297) --- CMakeLists.txt | 11 ++ cmake/backend_extensions.h.in | 2 + docs/ggml_extensions.md | 18 ++ docs/int8_convrot.md | 1 + ggml | 2 +- scripts/format-code.ps1 | 2 + scripts/format-code.sh | 1 + src/core/metal/CMakeLists.txt | 8 + src/core/metal/int8.metal | 143 +++++++++++++++ src/core/metal/int8.mm | 298 ++++++++++++++++++++++++++++++++ src/core/metal/int8_mpp.metal | 65 +++++++ src/core/metal/int8_source.h.in | 7 + src/model_manager_prefetch.cpp | 5 +- 13 files changed, 561 insertions(+), 2 deletions(-) create mode 100644 src/core/metal/CMakeLists.txt create mode 100644 src/core/metal/int8.metal create mode 100644 src/core/metal/int8.mm create mode 100644 src/core/metal/int8_mpp.metal create mode 100644 src/core/metal/int8_source.h.in diff --git a/CMakeLists.txt b/CMakeLists.txt index 8c5866d494..3e210379ae 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -111,6 +111,12 @@ if(SD_CUDA) endif() if(SD_METAL) + if(NOT SD_USE_UPSTREAM_GGML) + if(CMAKE_VERSION VERSION_LESS 3.16) + message(FATAL_ERROR "SD_METAL with SD_USE_UPSTREAM_GGML=OFF requires CMake 3.16 or later for OBJC and OBJCXX support") + endif() + enable_language(OBJC OBJCXX) + endif() message("-- Use Metal as backend stable-diffusion") set(GGML_METAL ON) endif() @@ -364,6 +370,11 @@ if (SD_CUDA) endif() endif() set(SD_BACKEND_REGISTRATIONS) +if (SD_METAL AND NOT SD_USE_UPSTREAM_GGML) + add_subdirectory(src/core/metal) + target_link_libraries(${SD_LIB} PRIVATE "-framework Foundation" "-framework Metal") + list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_metal_i8) +endif() if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA)) if(SD_HIPBLAS) include(cmake/hip.cmake) diff --git a/cmake/backend_extensions.h.in b/cmake/backend_extensions.h.in index 2dc1e0f6cd..602c3b4ae5 100644 --- a/cmake/backend_extensions.h.in +++ b/cmake/backend_extensions.h.in @@ -3,6 +3,8 @@ #include "core/cuda/registration.h" +void sd_ggml_register_metal_i8(ggml_backend_dev_t device); + void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device); void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device); void sd_ggml_register_vulkan_rms_affine(ggml_backend_dev_t device); diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 580ef58984..ea806fb88e 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -130,6 +130,9 @@ Services are discovered through `ggml_backend_reg_get_proc_address`: | `ggml_ext_gpu_device_v1` | Query CUDA/HIP/MUSA runtime kind, physical ordinal, shared-memory limit, warp size, and INT8 BLAS capability. | | `ggml_ext_gpu_context_v1` | Borrow the matching runtime's device and execution stream. | | `ggml_ext_gpu_blas_v1` | Borrow the matching cuBLAS, hipBLAS, or muBLAS handle on that stream. | +| `ggml_ext_metal_device_v1` | Borrow the Metal device for application pipeline compilation. | +| `ggml_ext_metal_buffer_v1` | Resolve a backend buffer range to its Metal buffer and offset, including private storage and split allocations. | +| `ggml_ext_metal_submit_v1` | Encode application commands on a backend-owned command buffer, ordered with native Metal work and tracked by backend synchronization. | The borrowed cuBLAS handle remains owned by the backend. Extensions must not retain or destroy it, or change its stream, pointer mode, math mode, or workspace. This @@ -161,6 +164,21 @@ CPU kernels and fusion rules live in `src/core/cpu/`, with registration collecte in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their respective backend directories. +`src/core/metal/int8.mm` registers Metal INT8 linear and ConvRot kernels. Their +Metal source is embedded at build time and compiled for the selected device +during registration. The public `ggml-ext-metal.h` services expose borrowed +objects without private backend headers. The submit callback only encodes work; +GGML owns submission, cancellation checks, command-buffer lifetime, and +asynchronous error tracking. Rotation scratch and any overlapping output use +the scheduler workspace. No full floating-point weight copy is allocated. +When runtime compilation supports Metal 4 TensorOps, INT8 multiplication uses +`matmul2d` with INT32 cooperative results and fused scale/bias restoration. +The SIMD-group implementation remains the fallback for older systems and +devices where the TensorOps pipeline is unavailable. +TensorOps multiplies at most 1024 pairs per partial result and combines those +results with wrapping integer addition. This preserves exact INT32 semantics +even when TensorOps uses floating-point arithmetic internally. + `ggml_ext_modulate(ctx, x, shift, scale)` emits `sd.modulate.v1` for F32 inputs when the runner context uses CPU, CUDA/HIP/MUSA, or Vulkan. Its explicit inputs are `x`, `scale`, and optional `shift`; it computes diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index d7948f6f39..c849dda14e 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -93,6 +93,7 @@ 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. +- Metal supports INT8 tensorwise and convrot on Apple silicon (Apple GPU family 7 or newer). It keeps weights in INT8 and rotates and quantizes activations on the GPU. On macOS/iOS 26 or newer, it probes Metal TensorOps `matmul2d` with INT8 inputs and INT32 cooperative results. It combines 1024-product partial sums using wrapping integer addition to preserve exact accumulation on older GPUs, then applies scales and bias in registers. Older systems or unavailable TensorOps pipelines use SIMD-group matrix multiplication: INT8 tiles are converted to FP16, and FP32 partial sums cover at most 256 products before being combined as integers, preserving the I32 accumulation contract. Convrot groups 1, 4, 16, 64, and 256 are supported. These kernels require the patched GGML Metal extension services and are enabled by `SD_METAL=ON`. - NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition. - AMD HIP devices in the CDNA, RDNA3 (including RDNA3.5), and RDNA4 families use the same INT8/H256 kernels with hipBLAS for I8 x I8 to I32 GEMM. Other AMD architectures fall back to CPU for INT8 matrix multiplication. - MUSA shares the CUDA-style H256 activation kernels. INT8 GEMM falls back to CPU because the MUSA backend does not advertise INT8 BLAS support. diff --git a/ggml b/ggml index a70d9e12b3..1e5639ec8f 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit a70d9e12b381f778bbb92505ecbe370469b43a52 +Subproject commit 1e5639ec8f3aa05c17ecddb039b86a09134ff03c diff --git a/scripts/format-code.ps1 b/scripts/format-code.ps1 index 479433b3f1..02097e6bb9 100644 --- a/scripts/format-code.ps1 +++ b/scripts/format-code.ps1 @@ -10,6 +10,8 @@ $patterns = @( "src/core/*.hpp" "src/core/cpu/*.cpp" "src/core/cpu/*.h" + "src/core/metal/*.mm" + "src/core/metal/*.metal" "src/core/cuda/*.cu" "src/core/cuda/*.cuh" "src/core/cuda/*.h" diff --git a/scripts/format-code.sh b/scripts/format-code.sh index 29b3313a9d..dd4d0fa95b 100644 --- a/scripts/format-code.sh +++ b/scripts/format-code.sh @@ -9,6 +9,7 @@ for f in src/*.cpp src/*.h src/*.hpp \ src/conditioning/*.cpp src/conditioning/*.h src/conditioning/*.hpp \ src/core/*.cpp src/core/*.h src/core/*.hpp \ src/core/cpu/*.cpp src/core/cpu/*.h \ + src/core/metal/*.mm src/core/metal/*.metal \ src/core/cuda/*.cu src/core/cuda/*.cuh src/core/cuda/*.h \ src/core/cuda/sageattention/*.cuh \ src/core/probe/*.cpp src/core/probe/*.h \ diff --git a/src/core/metal/CMakeLists.txt b/src/core/metal/CMakeLists.txt new file mode 100644 index 0000000000..9ff35d8b68 --- /dev/null +++ b/src/core/metal/CMakeLists.txt @@ -0,0 +1,8 @@ +set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS int8.metal int8_mpp.metal) +file(READ "${CMAKE_CURRENT_SOURCE_DIR}/int8.metal" SD_METAL_I8_SOURCE) +file(READ "${CMAKE_CURRENT_SOURCE_DIR}/int8_mpp.metal" SD_METAL_I8_MPP_SOURCE) +configure_file(int8_source.h.in "${CMAKE_CURRENT_BINARY_DIR}/int8_source.h" @ONLY) +target_sources(${SD_LIB} PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/int8.mm") +target_compile_options(${SD_LIB} PRIVATE "$<$:-fno-objc-arc>") +target_include_directories(${SD_LIB} PRIVATE "${CMAKE_CURRENT_BINARY_DIR}") +set_target_properties(${SD_LIB} PROPERTIES OBJCXX_STANDARD 17 OBJCXX_STANDARD_REQUIRED ON) diff --git a/src/core/metal/int8.metal b/src/core/metal/int8.metal new file mode 100644 index 0000000000..41f4a32b2b --- /dev/null +++ b/src/core/metal/int8.metal @@ -0,0 +1,143 @@ +#include +#include +using namespace metal; + +struct QuantParams { + uint k; + uint rows; + uint group; + uint unused; + ulong scale_offset; +}; + +kernel void sd_i8_rotate(device const float* input [[buffer(0)]], + device float* output [[buffer(1)]], + constant QuantParams& p [[buffer(2)]], + uint block [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]]) { + threadgroup float values[256]; + const ulong base = ulong(block) * p.group; + values[tid] = input[base + tid] * rsqrt(float(p.group)); + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint stride = 1; stride < p.group; stride *= 4) { + uint index = tid / (4 * stride) * (4 * stride) + tid % stride; + float a = values[index], b = values[index + stride]; + float c = values[index + 2 * stride], d = values[index + 3 * stride]; + float value; + switch ((tid / stride) % 4) { + case 0: + value = a + b + c - d; + break; + case 1: + value = a + b - c + d; + break; + case 2: + value = a - b + c + d; + break; + default: + value = -a + b + c + d; + break; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + values[tid] = value; + threadgroup_barrier(mem_flags::mem_threadgroup); + } + output[base + tid] = values[tid]; +} + +kernel void sd_i8_quantize(device const float* input [[buffer(0)]], + device char* output [[buffer(1)]], + constant QuantParams& p [[buffer(2)]], + uint row [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint lane [[thread_index_in_simdgroup]], + uint simd [[simdgroup_index_in_threadgroup]]) { + threadgroup float maxima[8]; + const ulong base = ulong(row) * p.k; + float amax = 0; + for (uint i = tid; i < p.k; i += 256) { + amax = max(amax, abs(input[base + i])); + } + amax = simd_max(amax); + if (lane == 0) + maxima[simd] = amax; + threadgroup_barrier(mem_flags::mem_threadgroup); + amax = simd_max(lane < 8 ? maxima[lane] : 0.f); + float scale = amax / 127.f; + float inverse = scale == 0 ? 0 : 1.f / scale; + for (uint i = tid; i < p.k; i += 256) { + output[base + i] = char(clamp(rint(input[base + i] * inverse), -127.f, 127.f)); + } + if (tid == 0) + ((device float*)(output + p.scale_offset))[row] = scale; +} + +struct MatParams { + uint k; + uint n; + uint rows; + uint flags; +}; + +kernel void sd_i8_multiply(device const char* weight [[buffer(0)]], + device const char* input [[buffer(1)]], + device const float* xs [[buffer(2)]], + device const float* ws [[buffer(3)]], + device const float* bias [[buffer(4)]], + device float* output [[buffer(5)]], + constant MatParams& p [[buffer(6)]], + uint2 tile [[threadgroup_position_in_grid]], + uint tid [[thread_index_in_threadgroup]], + uint simd [[simdgroup_index_in_threadgroup]]) { + threadgroup half a[32 * 32]; + threadgroup half b[32 * 32]; + threadgroup float partial[32 * 32]; + uint sums[8] = {}; + const uint row0 = tile.y * 32, col0 = tile.x * 32; + const uint sr = (simd / 2) * 16, sc = (simd % 2) * 16; + + // Half represents every INT8 value exactly. Each 256-term dot product fits + // in FP32's exact integer range; combine chunks as integers to avoid drift. + for (uint chunk = 0; chunk < p.k; chunk += 256) { + simdgroup_float8x8 acc[4]; + for (uint i = 0; i < 4; ++i) + acc[i] = make_filled_simdgroup_matrix(0.f); + for (uint kk = chunk; kk < min(chunk + 256, p.k); kk += 32) { + for (uint index = tid; index < 1024; index += 128) { + uint r = index / 32, k = index % 32; + a[index] = row0 + r < p.rows && kk + k < p.k ? half(input[ulong(row0 + r) * p.k + kk + k]) : half(0); + b[index] = col0 + r < p.n && kk + k < p.k ? half(weight[ulong(col0 + r) * p.k + kk + k]) : half(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint k = 0; k < 32; k += 8) { + simdgroup_half8x8 ma[2], mb[2]; + for (uint i = 0; i < 2; ++i) { + simdgroup_load(ma[i], a + (sr + i * 8) * 32 + k, 32); + simdgroup_load(mb[i], b + (sc + i * 8) * 32 + k, 32, ulong2(0), true); + } + for (uint i = 0; i < 4; ++i) { + simdgroup_multiply_accumulate(acc[i], ma[i / 2], mb[i % 2], acc[i]); + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i = 0; i < 4; ++i) { + simdgroup_store(acc[i], partial + (sr + (i / 2) * 8) * 32 + sc + (i % 2) * 8, 32); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + for (uint i = 0; i < 8; ++i) + sums[i] += uint(int(partial[tid + 128 * i])); + threadgroup_barrier(mem_flags::mem_threadgroup); + } + for (uint i = 0; i < 8; ++i) { + uint index = tid + 128 * i, row = row0 + index / 32, col = col0 + index % 32; + if (row < p.rows && col < p.n) { + float value = float(as_type(sums[i])) * xs[row]; + if (p.flags & 1) + value *= ws[col]; + if (p.flags & 2) + value += bias[col]; + output[ulong(row) * p.n + col] = value; + } + } +} diff --git a/src/core/metal/int8.mm b/src/core/metal/int8.mm new file mode 100644 index 0000000000..1fb2d568e4 --- /dev/null +++ b/src/core/metal/int8.mm @@ -0,0 +1,298 @@ +#include "core/ggml_extensions.h" +#include "ggml-ext-metal.h" +#include "int8_source.h" + +#import +#import + +#include +#include +#include + +namespace sd_metal { + + struct Kernels { + id rotate = nil; + id quantize = nil; + id multiply = nil; + id multiply_mpp = nil; + ggml_ext_metal_buffer_fn buffer = nullptr; + ggml_ext_metal_submit_fn submit = nullptr; + + ~Kernels() { + [rotate release]; + [quantize release]; + [multiply release]; + [multiply_mpp release]; + } + }; + + static id tensor_multiply(id device) { + if (@available(macOS 26.0, iOS 26.0, *)) { + MTLCompileOptions* options = [[MTLCompileOptions alloc] init]; + // Keep building with older SDKs; the optional shader is compiled at runtime. + options.languageVersion = static_cast(4 << 16); + options.fastMathEnabled = NO; + NSError* error = nil; + id library = [device newLibraryWithSource:[NSString stringWithUTF8String:sd_metal_i8_mpp_source] options:options error:&error]; + [options release]; + id function = [library newFunctionWithName:@"sd_i8_multiply_mpp"]; + id pipeline = function ? [device newComputePipelineStateWithFunction:function error:&error] : nil; + [function release]; + [library release]; + if (pipeline && pipeline.threadExecutionWidth == 32 && pipeline.maxTotalThreadsPerThreadgroup >= 128) { + return pipeline; + } + [pipeline release]; + NSLog(@"sd.cpp Metal INT8 TensorOps unavailable; using SIMD-group fallback: %@", error); + } + return nil; + } + + struct QuantParams { + uint32_t k, rows, group, unused; + uint64_t scale_offset; + }; + + struct MatParams { + uint32_t k, n, rows, flags; + }; + + struct Execution { + const ggml_ext_execution_context* context; + ggml_tensor* node; + Kernels* kernels; + SDI8Operation op = {}; + int group = 0; + bool linear = false; + }; + + static bool dimensions(const ggml_tensor* node) { + return node->ne[0] <= INT_MAX && ggml_nrows(node) <= INT_MAX; + } + + static bool supports_i8(ggml_backend_dev_t, const ggml_tensor* node, void*) { + SDI8Operation op; + return sd_ggml_i8_describe(node, &op) && dimensions(node) && dimensions(op.weight) && + uint64_t(op.weight->ne[0]) * ggml_nrows(node) <= UINT32_MAX; + } + + static bool supports_convrot(ggml_backend_dev_t, const ggml_tensor* node, void*) { + int group; + return sd_ggml_convrot_describe(node, &group) && dimensions(node->src[0]) && + uint64_t(ggml_nelements(node->src[0])) <= UINT32_MAX; + } + + static size_t aligned(size_t bytes) { + return (bytes + 255) / 256 * 256; + } + + static size_t packed_size(int64_t k, int64_t rows) { + return aligned(size_t(k) * ((rows + 3) / 4 * 4) + rows * sizeof(float)); + } + + static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) { + auto* node = nodes[0]; + SDI8Operation op; + const bool linear = sd_ggml_i8_describe(node, &op); + const auto* input = linear ? op.input : node->src[0]; + const int64_t rows = linear ? ggml_nrows(node) : ggml_nrows(input); + int group; + if (linear) { + group = op.group_size; + } else if (!sd_ggml_convrot_describe(node, &group)) { + return GGML_STATUS_FAILED; + } + *bytes = aligned(sd_ggml_modulate_workspace_size(node)); + if (input->type == GGML_TYPE_F32) { + if (group > 1) + *bytes += aligned(ggml_nbytes(input)); + if (linear) + *bytes += packed_size(input->ne[0], rows); + } + return GGML_STATUS_SUCCESS; + } + + static bool bind_tensor(Kernels* kernels, const ggml_tensor* tensor, ggml_ext_metal_buffer& result) { + auto buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (!buffer || !tensor->data) + return false; + const size_t offset = uintptr_t(tensor->data) - uintptr_t(ggml_backend_buffer_get_base(buffer)); + return kernels->buffer(buffer, offset, ggml_nbytes(tensor), &result); + } + + static void set_buffer(id encoder, const ggml_ext_metal_buffer& buffer, int index) { + [encoder setBuffer:(id)buffer.buffer offset:buffer.offset atIndex:index]; + } + + static ggml_status encode(void* command_buffer, void* user_data) { + auto& e = *static_cast(user_data); + auto* kernels = e.kernels; + const auto* input = e.linear ? e.op.input : e.node->src[0]; + const uint32_t k = uint32_t(input->ne[0]); + const uint32_t rows = uint32_t(e.linear ? ggml_nrows(e.node) : ggml_nrows(input)); + id command = (id)command_buffer; + ggml_ext_metal_buffer x, dst, output, rotated, packed, xs, w, ws, bias; + if (!bind_tensor(kernels, input, x) || !bind_tensor(kernels, e.node, dst)) + return GGML_STATUS_FAILED; + output = dst; + size_t offset = 0; + auto scratch = [&](size_t bytes, ggml_ext_metal_buffer& binding) { + bool ok = kernels->buffer(e.context->workspace, offset, bytes, &binding); + offset += aligned(bytes); + return ok; + }; + const size_t output_scratch = sd_ggml_modulate_workspace_size(e.node); + if (output_scratch && !scratch(output_scratch, output)) + return GGML_STATUS_FAILED; + + if (input->type == GGML_TYPE_F32) { + rotated = x; + if (e.group > 1 && !scratch(ggml_nbytes(input), rotated)) + return GGML_STATUS_FAILED; + packed = output; + if (e.linear && !scratch(packed_size(k, rows), packed)) + return GGML_STATUS_FAILED; + QuantParams p = {k, rows, uint32_t(e.group), 0, uint64_t(k) * ((uint64_t(rows) + 3) / 4 * 4)}; + id clear = [command blitCommandEncoder]; + [clear fillBuffer:(id)packed.buffer range:NSMakeRange(packed.offset, e.linear ? packed_size(k, rows) : ggml_nbytes(e.node)) value:0]; + [clear endEncoding]; + if (e.group > 1) { + id encoder = [command computeCommandEncoder]; + [encoder setComputePipelineState:kernels->rotate]; + set_buffer(encoder, x, 0); + set_buffer(encoder, rotated, 1); + [encoder setBytes:&p length:sizeof(p) atIndex:2]; + [encoder dispatchThreadgroups:MTLSizeMake(uint64_t(k) * rows / e.group, 1, 1) + threadsPerThreadgroup:MTLSizeMake(e.group, 1, 1)]; + [encoder endEncoding]; + } + id encoder = [command computeCommandEncoder]; + [encoder setComputePipelineState:kernels->quantize]; + set_buffer(encoder, rotated, 0); + set_buffer(encoder, packed, 1); + [encoder setBytes:&p length:sizeof(p) atIndex:2]; + [encoder dispatchThreadgroups:MTLSizeMake(rows, 1, 1) threadsPerThreadgroup:MTLSizeMake(256, 1, 1)]; + [encoder endEncoding]; + x = packed; + xs = packed; + xs.offset += p.scale_offset; + } else if (e.op.packed) { + xs = x; + xs.offset += size_t(k) * ((uint64_t(rows) + 3) / 4 * 4); + } else if (!bind_tensor(kernels, e.op.input_scale, xs)) { + return GGML_STATUS_FAILED; + } + + if (e.linear) { + if (!bind_tensor(kernels, e.op.weight, w)) + return GGML_STATUS_FAILED; + ws = bias = xs; + if (e.op.weight_scale && !bind_tensor(kernels, e.op.weight_scale, ws)) + return GGML_STATUS_FAILED; + if (e.op.bias && !bind_tensor(kernels, e.op.bias, bias)) + return GGML_STATUS_FAILED; + MatParams p = {k, uint32_t(e.node->ne[0]), rows, uint32_t((e.op.weight_scale ? 1 : 0) | (e.op.bias ? 2 : 0))}; + id encoder = [command computeCommandEncoder]; + [encoder setComputePipelineState:kernels->multiply_mpp ? kernels->multiply_mpp : kernels->multiply]; + set_buffer(encoder, w, 0); + set_buffer(encoder, x, 1); + set_buffer(encoder, xs, 2); + set_buffer(encoder, ws, 3); + set_buffer(encoder, bias, 4); + set_buffer(encoder, output, 5); + [encoder setBytes:&p length:sizeof(p) atIndex:6]; + [encoder dispatchThreadgroups:MTLSizeMake((p.n + 31) / 32, (rows + 31) / 32, 1) + threadsPerThreadgroup:MTLSizeMake(128, 1, 1)]; + [encoder endEncoding]; + } + if (output_scratch) { + id copy = [command blitCommandEncoder]; + [copy copyFromBuffer:(id)output.buffer + sourceOffset:output.offset + toBuffer:(id)dst.buffer + destinationOffset:dst.offset + size:ggml_nbytes(e.node)]; + [copy endEncoding]; + } + return GGML_STATUS_SUCCESS; + } + + static ggml_status execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void* user_data) { + Execution e = {context, nodes[0], static_cast(user_data)}; + e.linear = sd_ggml_i8_describe(e.node, &e.op); + if (e.linear) { + e.group = e.op.group_size; + } else if (!sd_ggml_convrot_describe(e.node, &e.group)) { + return GGML_STATUS_FAILED; + } + return e.kernels->submit(context->backend, encode, &e); + } + +} + +void sd_ggml_register_metal_i8(ggml_backend_dev_t device) { + using namespace sd_metal; + auto reg = ggml_backend_dev_backend_reg(device); + auto get_device = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_metal_device_v1")); + auto get_buffer = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_metal_buffer_v1")); + auto submit = reinterpret_cast(ggml_backend_reg_get_proc_address(reg, "ggml_ext_metal_submit_v1")); + if (!get_device || !get_buffer || !submit) + return; + @autoreleasepool { + id metal = (id)get_device(device); + if (!metal || ![metal supportsFamily:MTLGPUFamilyApple7]) + return; + auto kernels = std::make_unique(); + kernels->buffer = get_buffer; + kernels->submit = submit; + MTLCompileOptions* options = [[MTLCompileOptions alloc] init]; + options.fastMathEnabled = NO; + NSError* error = nil; + id library = [metal newLibraryWithSource:[NSString stringWithUTF8String:sd_metal_i8_source] options:options error:&error]; + [options release]; + if (!library) { + NSLog(@"sd.cpp Metal INT8 compilation failed: %@", error); + return; + } + auto pipeline = [&](NSString* name) { + id function = [library newFunctionWithName:name]; + id result = function ? [metal newComputePipelineStateWithFunction:function error:&error] : nil; + [function release]; + return result; + }; + kernels->rotate = pipeline(@"sd_i8_rotate"); + kernels->quantize = pipeline(@"sd_i8_quantize"); + kernels->multiply = pipeline(@"sd_i8_multiply"); + [library release]; + if (!kernels->rotate || !kernels->quantize || !kernels->multiply || + kernels->quantize.threadExecutionWidth != 32 || kernels->multiply.threadExecutionWidth != 32 || + kernels->rotate.maxTotalThreadsPerThreadgroup < 256 || + kernels->quantize.maxTotalThreadsPerThreadgroup < 256 || kernels->multiply.maxTotalThreadsPerThreadgroup < 128) { + NSLog(@"sd.cpp Metal INT8 pipelines unavailable: %@", error); + return; + } + kernels->multiply_mpp = tensor_multiply(metal); + if (kernels->multiply_mpp) { + NSLog(@"sd.cpp Metal INT8: using TensorOps INT8 x INT8 -> INT32"); + } + ggml_ext_kernel_desc desc = {}; + desc.abi_version = GGML_EXT_ABI_VERSION; + desc.name = "sd.linear_i8.metal.v1"; + desc.op = GGML_OP_CUSTOM; + desc.custom_op = sd_ggml_i8_op(); + desc.priority = 100; + desc.supports = supports_i8; + desc.prepare = prepare; + desc.destroy = [](void*, void*) {}; + desc.execute = execute; + desc.user_data = kernels.get(); + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + desc.name = "sd.quantize_i8_convrot.metal.v1"; + desc.custom_op = sd_ggml_convrot_op(); + desc.supports = supports_convrot; + GGML_ASSERT(ggml_ext_register_kernel(device, &desc)); + static std::vector> registrations; + registrations.push_back(std::move(kernels)); + } +} diff --git a/src/core/metal/int8_mpp.metal b/src/core/metal/int8_mpp.metal new file mode 100644 index 0000000000..af2a9ea0b4 --- /dev/null +++ b/src/core/metal/int8_mpp.metal @@ -0,0 +1,65 @@ +#include +#include +#include +using namespace metal; +using namespace mpp::tensor_ops; + +struct MatParams { + uint k; + uint n; + uint rows; + uint flags; +}; + +// TensorOps requires explicit signed INT8 and non-const tensor element types. +kernel void sd_i8_multiply_mpp(device int8_t* weight [[buffer(0)]], + device int8_t* input [[buffer(1)]], + device const float* xs [[buffer(2)]], + device const float* ws [[buffer(3)]], + device const float* bias [[buffer(4)]], + device float* output [[buffer(5)]], + constant MatParams& p [[buffer(6)]], + uint2 tile [[threadgroup_position_in_grid]]) { + auto a = tensor, tensor_inline>(input, dextents(p.k, p.rows)); + auto b = tensor, tensor_inline>(weight, dextents(p.k, p.n)); + constexpr int chunk_size = 1024; + constexpr auto descriptor = matmul2d_descriptor(32, 32, chunk_size, false, true, false); + matmul2d> multiply; + auto left = a.slice(0, tile.y * 32); + auto right = b.slice(0, tile.x * 32); + auto acc = multiply.get_destination_cooperative_tensor(); + auto sums = multiply.get_destination_cooperative_tensor(); +#pragma unroll + for (uint i = 0; i < sums.get_capacity(); ++i) { + if (sums.is_valid_element(i)) + sums[i] = 0; + } + // TensorOps can round long integer dot products on older GPUs. Bound each + // partial sum to FP32's exact integer range, then combine with wrapping + // integer addition outside TensorOps. + for (uint k = 0; k < p.k; k += chunk_size) { + left = a.slice(k, tile.y * 32); + right = b.slice(k, tile.x * 32); + multiply.run(left, right, acc); +#pragma unroll + for (uint i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) + sums[i] = as_type(as_type(sums[i]) + as_type(acc[i])); + } + } +#pragma unroll + for (uint i = 0; i < acc.get_capacity(); ++i) { + if (acc.is_valid_element(i)) { + auto position = acc.get_multidimensional_index(i); + uint col = tile.x * 32 + position[0], row = tile.y * 32 + position[1]; + if (col < p.n && row < p.rows) { + float value = float(sums[i]) * xs[row]; + if (p.flags & 1) + value *= ws[col]; + if (p.flags & 2) + value += bias[col]; + output[ulong(row) * p.n + col] = value; + } + } + } +} diff --git a/src/core/metal/int8_source.h.in b/src/core/metal/int8_source.h.in new file mode 100644 index 0000000000..13a07c6c97 --- /dev/null +++ b/src/core/metal/int8_source.h.in @@ -0,0 +1,7 @@ +#ifndef __SD_CORE_METAL_INT8_SOURCE_H__ +#define __SD_CORE_METAL_INT8_SOURCE_H__ + +static const char sd_metal_i8_source[] = R"SD_METAL(@SD_METAL_I8_SOURCE@)SD_METAL"; +static const char sd_metal_i8_mpp_source[] = R"SD_METAL(@SD_METAL_I8_MPP_SOURCE@)SD_METAL"; + +#endif diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp index 4056b8f0ea..128cae237a 100644 --- a/src/model_manager_prefetch.cpp +++ b/src/model_manager_prefetch.cpp @@ -33,7 +33,10 @@ 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.submitted && block.transfer_backend != nullptr) { + } + // Event completion does not release backend-owned transfer resources, such + // as Metal's temporary host mappings and command buffers. + if (block.submitted && block.transfer_backend != nullptr) { ggml_backend_synchronize(block.transfer_backend); } block.transfer_backend = nullptr; From 80c4707c535f877f1061d212dcf04bcfe41e35d8 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 00:43:02 +0800 Subject: [PATCH 33/57] feat: support HyperFlow H3 Flash (#2299) Co-authored-by: Video Rebirth --- docs/minimax_h3.md | 15 ++ examples/common/common.cpp | 10 +- examples/server/recipes/defaults.md | 2 + .../recipes/hyperflow-h3-flash-fl2va.json | 75 ++++++ .../recipes/hyperflow-h3-flash-ref2va.json | 74 ++++++ src/model/diffusion/minimax_h3.hpp | 231 ++++++++++++++---- src/model/diffusion/model.hpp | 1 + src/model_manager_prefetch.cpp | 5 +- src/pipeline/diffusion_engine.cpp | 6 +- 9 files changed, 370 insertions(+), 49 deletions(-) create mode 100644 examples/server/recipes/hyperflow-h3-flash-fl2va.json create mode 100644 examples/server/recipes/hyperflow-h3-flash-ref2va.json diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index bc1f19c770..b5514db69d 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -169,3 +169,18 @@ Notes for this mode: block's working memory must also fit within the VRAM budget. - The default video flow shift is 12. The audio stream is mapped internally to its shift of 3, so the regular samplers can operate on the packed AV latent. + +## HyperFlow H3 Flash + +[HyperFlow H3 Flash](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) +uses the shared components and input options described in this guide. Replace +`--diffusion-model` with the matching HyperFlow FL2VA or Ref2VA checkpoint and +use these sampling options: + +```sh +--sampling-method euler --cfg-scale 1 --sigmas hyperflow-h3-flash-6 +``` + +Use `hyperflow-h3-flash-4` or `hyperflow-h3-flash-8` for 4 or 8 steps. These +presets also work with `--hires-sigmas`; their length determines the step count. +They already include video flow shift 12, so keep the default flow shift. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 79349241cf..a3496f4fcc 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1540,6 +1540,12 @@ ArgOptions SDGenerationParams::get_options() { sigmas_str = "1.0,0.978453,0.95418,0.926626,0.89508,0.845148,0.704534,0.414568,0.0"; } else if (sigmas_str == "p-h3-image") { sigmas_str = "1.0,0.984753,0.967905,0.951688,0.925372,0.881923,0.810083,0.528718,0.0"; + } else if (sigmas_str == "hyperflow-h3-flash-4") { + sigmas_str = "1.0,0.962575853,0.755781889,0.312497854,0.0"; + } else if (sigmas_str == "hyperflow-h3-flash-6") { + sigmas_str = "1.0,0.982655942,0.914642572,0.674755692,0.385223746,0.159055337,0.0"; + } else if (sigmas_str == "hyperflow-h3-flash-8") { + sigmas_str = "1.0,0.993909776,0.9842875,0.966063678,0.923076928,0.834942341,0.696852028,0.468753338,0.0"; } if (!sigmas_str.empty() && sigmas_str.front() == '[') { sigmas_str.erase(0, 1); @@ -1787,11 +1793,11 @@ ArgOptions SDGenerationParams::get_options() { on_scheduler_arg}, {"", "--sigmas", - "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image.", + "custom sigma values for the sampler, comma-separated (e.g., \"14.61,7.8,3.5,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.", on_sigmas_arg}, {"", "--hires-sigmas", - "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image.", + "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.", on_hires_sigmas_arg}, {"", "--skip-layers", diff --git a/examples/server/recipes/defaults.md b/examples/server/recipes/defaults.md index efe596c92b..bd0abd5636 100644 --- a/examples/server/recipes/defaults.md +++ b/examples/server/recipes/defaults.md @@ -42,6 +42,8 @@ For community edit adapters and other optional extensions, these are the base-mo | `flux3-klein` | 1 | 16 | 1024 Γ— 1024 | [Source](https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/README.md) | | `flux3-klein-base` | 4 | 20 | 1024 Γ— 1024 | [Source](https://raw.githubusercontent.com/black-forest-labs/flux-3-image-klein/main/src/flux3/pipeline.py) | | `hidream-o1-dev` | 1 | 20 | 1024 Γ— 1024 | [Source](https://huggingface.co/HiDream-ai/HiDream-O1-Image-Dev) | +| `hyperflow-h3-flash-fl2va` | 1 | 6 | 1344 Γ— 768 | [Source](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) Two-time FL2VA checkpoint; native text and first/last-frame generation. Fixed 6-step video sigma grid with shift 12, 124 frames at 24 FPS. Reuses MiniMax-H3 components. | +| `hyperflow-h3-flash-ref2va` | 1 | 6 | 1344 Γ— 768 | [Source](https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash) Separate two-time Ref2VA checkpoint; image references. Same schedule and components as FL2VA; video/audio references remain unverified. | | `hunyuan-video-15-i2v` | 6 | 20 | 1280 Γ— 720 | [Source](https://huggingface.co/tencent/HunyuanVideo-1.5) 720p base I2V defaults; CFG/step-distilled and 480p checkpoints use different settings. | | `hunyuan-video-15-t2v` | 6 | 20 | 1280 Γ— 720 | [Source](https://huggingface.co/tencent/HunyuanVideo-1.5) 720p base T2V defaults; CFG-distilled / 480p checkpoints need their published settings. | | `ideogram4` | 7 | 20 | 1024 Γ— 1024 | [Source](https://raw.githubusercontent.com/ideogram-oss/ideogram4/main/docs/inference.md) Constant-guidance pipeline defaults; this recipe does not reproduce the CLI’s per-step guidance schedule. | diff --git a/examples/server/recipes/hyperflow-h3-flash-fl2va.json b/examples/server/recipes/hyperflow-h3-flash-fl2va.json new file mode 100644 index 0000000000..64e79d5436 --- /dev/null +++ b/examples/server/recipes/hyperflow-h3-flash-fl2va.json @@ -0,0 +1,75 @@ +{ + "schema_version": 1, + "id": "hyperflow-h3-flash-fl2va", + "name": "HyperFlow H3 Flash FL2VA", + "size_multiple": 32, + "description": "Generate video and audio from text, optionally with start/end images. Use the HyperFlow H3 Flash FL2VA transformer. Reuses MiniMax-H3 components with the HyperFlow 6-step schedule. Audio VAE is optional for silent video.", + "model_url": "https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash", + "slots": [ + { + "id": "diffusion_model", + "required": true, + "label": "HyperFlow H3 Flash FL2VA Transformer" + }, + { + "id": "llm", + "required": true, + "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder" + }, + { + "id": "vae", + "required": false, + "label": "MiniMax H3 Video VAE" + }, + { + "id": "tae", + "required": false + }, + { + "id": "audio_vae", + "required": false, + "label": "MiniMax H3 Audio VAE" + } + ], + "required_any": [ + [ + "vae", + "tae" + ] + ], + "video_constraints": { + "frame_step": 17, + "frame_offset": 5, + "frame_rounding": "ceil", + "min_frames": 5, + "fixed_fps": 24 + }, + "defaults": { + "video_frames": 124, + "fps": 24, + "cfg_scale": 1, + "width": 1344, + "height": 768, + "sampling_method": "euler", + "steps": 6, + "flow_shift": 12, + "custom_sigmas": [ + 1.0, + 0.982655942, + 0.914642572, + 0.674755692, + 0.385223746, + 0.159055337, + 0.0 + ] + }, + "features": [ + "end_image", + "flow_shift", + "init_image" + ], + "disabled_features": [ + "img2img", + "mask_image" + ] +} diff --git a/examples/server/recipes/hyperflow-h3-flash-ref2va.json b/examples/server/recipes/hyperflow-h3-flash-ref2va.json new file mode 100644 index 0000000000..871a69063d --- /dev/null +++ b/examples/server/recipes/hyperflow-h3-flash-ref2va.json @@ -0,0 +1,74 @@ +{ + "schema_version": 1, + "id": "hyperflow-h3-flash-ref2va", + "name": "HyperFlow H3 Flash Ref2VA", + "size_multiple": 32, + "description": "Generate video and audio with reference images. Use the HyperFlow H3 Flash Ref2VA transformer; video/audio references are not validated for these weights. Reuses MiniMax-H3 components with the HyperFlow 6-step schedule. Audio VAE is optional for silent video.", + "model_url": "https://huggingface.co/VideoRebirthAI/hyperflow-h3-flash", + "slots": [ + { + "id": "diffusion_model", + "required": true, + "label": "HyperFlow H3 Flash REF2VA Transformer" + }, + { + "id": "llm", + "required": true, + "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder" + }, + { + "id": "vae", + "required": false, + "label": "MiniMax H3 Video VAE" + }, + { + "id": "tae", + "required": false + }, + { + "id": "audio_vae", + "required": false, + "label": "MiniMax H3 Audio VAE" + } + ], + "required_any": [ + [ + "vae", + "tae" + ] + ], + "video_constraints": { + "frame_step": 17, + "frame_offset": 5, + "frame_rounding": "ceil", + "min_frames": 5, + "fixed_fps": 24 + }, + "defaults": { + "video_frames": 124, + "fps": 24, + "cfg_scale": 1, + "width": 1344, + "height": 768, + "sampling_method": "euler", + "steps": 6, + "flow_shift": 12, + "custom_sigmas": [ + 1.0, + 0.982655942, + 0.914642572, + 0.674755692, + 0.385223746, + 0.159055337, + 0.0 + ] + }, + "features": [ + "flow_shift", + "ref_images" + ], + "disabled_features": [ + "img2img", + "mask_image" + ] +} diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 5f388d84f4..9a8d91388c 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -24,6 +24,19 @@ namespace MiniMaxH3 { constexpr float FRAME_RESCALE = 5.f / 3.f; constexpr float VISUAL_COND_TIMESTEP = 0.999f; + struct LowRankShape { + int64_t q = 0; + int64_t k = 0; + int64_t v = 0; + int64_t out = 0; + int64_t fc1 = 0; + int64_t fc2 = 0; + + bool splits_qkv() const { + return q > 0 || k > 0 || v > 0; + } + }; + struct Config { int64_t hidden_size = 5376; int64_t num_layers = 50; @@ -45,6 +58,8 @@ namespace MiniMaxH3 { float norm_eps = 1e-5f; float qk_norm_eps = 1e-5f; float final_norm_eps = 1e-5f; + bool two_time = false; + std::vector block_ranks; bool uses_adaln_curves() const { return adaln_curve_grid > 0; @@ -86,12 +101,36 @@ namespace MiniMaxH3 { if (const auto* weight = find("blocks.0.attn.q_norm.weight")) { config.attention_head_dim = weight->logical_ne0(); } + auto rank_of = [&](const std::string& name) -> int64_t { + const auto* down = find(name + ".down.weight"); + return down == nullptr ? 0 : down->ne[1]; + }; + config.block_ranks.resize(static_cast(config.num_layers)); + for (int64_t i = 0; i < config.num_layers; ++i) { + std::string block = "blocks." + std::to_string(i) + "."; + auto& ranks = config.block_ranks[static_cast(i)]; + ranks.q = rank_of(block + "attn.q_proj"); + ranks.k = rank_of(block + "attn.k_proj"); + ranks.v = rank_of(block + "attn.v_proj"); + ranks.out = rank_of(block + "attn.out_proj"); + ranks.fc1 = rank_of(block + "mlp.fc1"); + ranks.fc2 = rank_of(block + "mlp.fc2"); + } if (const auto* weight = find("blocks.0.attn.qkv_proj.weight")) { config.num_attention_heads = weight->ne[1] / (3 * config.attention_head_dim); + } else { + const auto* q_weight = find("blocks.0.attn.q_proj.up.weight"); + q_weight = q_weight != nullptr ? q_weight : find("blocks.0.attn.q_proj.weight"); + if (q_weight != nullptr) { + config.num_attention_heads = q_weight->ne[1] / config.attention_head_dim; + } } - if (const auto* weight = find("blocks.0.mlp.fc1.weight")) { - config.ffn_hidden_size = weight->ne[1] / 2; + const auto* fc1 = find("blocks.0.mlp.fc1.weight"); + fc1 = fc1 != nullptr ? fc1 : find("blocks.0.mlp.fc1.up.weight"); + if (fc1 != nullptr) { + config.ffn_hidden_size = fc1->ne[1] / 2; } + config.two_time = find("endpoint_time_embedder.proj_in.weight") != nullptr; if (const auto* weight = find("condition_proj.weight")) { config.text_dim = weight->logical_ne0(); } @@ -112,13 +151,14 @@ namespace MiniMaxH3 { } LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 - ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64, + ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 ", two_time=%d", config.num_layers, config.hidden_size, config.num_attention_heads, config.attention_head_dim, config.ffn_hidden_size, - config.adaln_curve_grid); + config.adaln_curve_grid, + config.two_time); return config; } }; @@ -141,31 +181,65 @@ namespace MiniMaxH3 { } }; + struct LowRankLinear : public UnaryBlock { + LowRankLinear(int64_t in_features, int64_t rank, int64_t out_features, bool force_prec_f32 = false, float scale = 1.f) { + blocks["down"] = std::make_shared(in_features, rank, false, false, force_prec_f32, scale); + blocks["up"] = std::make_shared(rank, out_features, false, false, force_prec_f32, scale); + } + + std::shared_ptr part(const char* name) { + return std::dynamic_pointer_cast(blocks[name]); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto down = std::dynamic_pointer_cast(blocks["down"]); + auto up = std::dynamic_pointer_cast(blocks["up"]); + return up->forward(ctx, down->forward(ctx, x)); + } + }; + + static std::shared_ptr make_linear(int64_t in_features, int64_t out_features, int64_t rank, bool force_prec_f32 = false, float scale = 1.f) { + if (rank > 0) { + return std::make_shared(in_features, rank, out_features, force_prec_f32, scale); + } + return std::make_shared(in_features, out_features, false, false, force_prec_f32, scale); + } + struct MLP : public UnaryBlock { int64_t ffn_hidden_size; MLP(int64_t hidden_size, - int64_t ffn_hidden_size) + int64_t ffn_hidden_size, + int64_t fc1_rank = 0, + int64_t fc2_rank = 0) : 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); + blocks["fc1"] = make_linear(hidden_size, ffn_hidden_size * 2, fc1_rank, true, 1.f / 128.f); + blocks["fc2"] = make_linear(ffn_hidden_size, hidden_size, fc2_rank, true, 1.f / 128.f); } 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]); + auto set_scale = [&](const std::shared_ptr& block, const std::string& name) { + auto linear = std::dynamic_pointer_cast(block); linear->set_scale(get_type(prefix + name + ".weight", tensors, GGML_TYPE_F32) == GGML_TYPE_I8 ? 1.f : 1.f / 128.f); + }; + for (const std::string name : {"fc1", "fc2"}) { + if (auto low_rank = std::dynamic_pointer_cast(blocks[name])) { + set_scale(low_rank->part("down"), name + ".down"); + set_scale(low_rank->part("up"), name + ".up"); + } else { + set_scale(blocks[name], name); + } } } 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 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); @@ -233,17 +307,35 @@ namespace MiniMaxH3 { int64_t heads; int64_t head_dim; float head_chunk_budget_gib = DEFAULT_ATTN_CHUNK_BUDGET_GIB; + bool split_projections; Attention(int64_t hidden_size, int64_t heads, int64_t head_dim, - float eps) - : heads(heads), head_dim(head_dim) { - int64_t inner = heads * head_dim; - blocks["qkv_proj"] = std::make_shared(hidden_size, inner * 3, false); + float eps, + const LowRankShape& ranks = {}) + : heads(heads), head_dim(head_dim), split_projections(ranks.splits_qkv()) { + int64_t inner = heads * head_dim; + if (split_projections) { + blocks["q_proj"] = make_linear(hidden_size, inner, ranks.q); + blocks["k_proj"] = make_linear(hidden_size, inner, ranks.k); + blocks["v_proj"] = make_linear(hidden_size, inner, ranks.v); + } else { + blocks["qkv_proj"] = std::make_shared(hidden_size, inner * 3, false); + } blocks["q_norm"] = std::make_shared(head_dim, eps); blocks["k_norm"] = std::make_shared(head_dim, eps); - blocks["out_proj"] = std::make_shared(inner, hidden_size, false); + blocks["out_proj"] = make_linear(inner, hidden_size, ranks.out); + } + + ggml_tensor* project_packed_qkv(GGMLRunnerContext* ctx, ggml_tensor* x) { + if (!split_projections) { + return std::dynamic_pointer_cast(blocks["qkv_proj"])->forward(ctx, x); + } + auto q = std::dynamic_pointer_cast(blocks["q_proj"])->forward(ctx, x); + auto k = std::dynamic_pointer_cast(blocks["k_proj"])->forward(ctx, x); + auto v = std::dynamic_pointer_cast(blocks["v_proj"])->forward(ctx, x); + return ggml_concat(ctx->ggml_ctx, ggml_concat(ctx->ggml_ctx, q, k, 0), v, 0); } std::vector split_qkv(GGMLRunnerContext* ctx, ggml_tensor* projected) { @@ -260,7 +352,6 @@ namespace MiniMaxH3 { std::vector project_qkv(GGMLRunnerContext* ctx, const std::vector& inputs, const std::function& prepare_input = {}) { - auto projection = std::dynamic_pointer_cast(blocks["qkv_proj"]); int64_t sequence = 0; for (auto input : inputs) { sequence += input->ne[1]; @@ -279,7 +370,7 @@ namespace MiniMaxH3 { ggml_tensor* projected = nullptr; for (size_t i = 0; i < inputs.size(); ++i) { auto input = prepare_input ? prepare_input(inputs[i], i) : inputs[i]; - auto part = projection->forward(ctx, input); + auto part = project_packed_qkv(ctx, input); projected = projected == nullptr ? part : ggml_concat(ctx->ggml_ctx, projected, part, 1); } return split_qkv(ctx, projected); @@ -298,7 +389,7 @@ namespace MiniMaxH3 { if (prepare_input) { part = prepare_input(part, i); } - auto qkv = split_qkv(ctx, projection->forward(ctx, part)); + auto qkv = split_qkv(ctx, project_packed_qkv(ctx, part)); for (int j = 0; j < 3; ++j) { chunks[j].push_back(ggml_cont(ctx->ggml_ctx, qkv[j])); } @@ -312,7 +403,7 @@ namespace MiniMaxH3 { } ggml_tensor* project_output(GGMLRunnerContext* ctx, ggml_tensor* x) { - return std::dynamic_pointer_cast(blocks["out_proj"])->forward(ctx, x); + return std::dynamic_pointer_cast(blocks["out_proj"])->forward(ctx, x); } ggml_tensor* attend(GGMLRunnerContext* ctx, @@ -489,16 +580,19 @@ namespace MiniMaxH3 { struct TransformerBlock : public GGMLBlock { Config config; - explicit TransformerBlock(const Config& config) + TransformerBlock(const Config& config, const LowRankShape& ranks) : config(config) { 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.qk_norm_eps, + ranks); blocks["mlp"] = std::make_shared(config.hidden_size, - config.ffn_hidden_size); + config.ffn_hidden_size, + ranks.fc1, + ranks.fc2); blocks["adaln_proj"] = std::make_shared(config.time_embed_dim, config.hidden_size, 6, @@ -602,10 +696,16 @@ namespace MiniMaxH3 { blocks["time_embedder"] = std::make_shared(config.timestep_input_dim, config.time_embed_hidden_size, config.time_embed_dim); + if (config.two_time) { + blocks["endpoint_time_embedder"] = std::make_shared(config.timestep_input_dim, + config.time_embed_hidden_size, + config.time_embed_dim); + } } blocks["token_refiner"] = std::make_shared(config); for (int64_t i = 0; i < config.num_layers; ++i) { - blocks["blocks." + std::to_string(i)] = std::make_shared(config); + blocks["blocks." + std::to_string(i)] = + std::make_shared(config, config.block_ranks.empty() ? LowRankShape{} : config.block_ranks[static_cast(i)]); } blocks["final_layer"] = std::make_shared(config); } @@ -615,6 +715,9 @@ namespace MiniMaxH3 { const std::string prefix = "") override { GGMLBlock::init_params(ctx, tensors, prefix); params["rope.inv_freq"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config.rope_inv_freq_len); + if (config.two_time) { + params["hyperflow_gate"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + } if (config.uses_adaln_curves()) { params["adaln_t_table"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, @@ -639,11 +742,20 @@ namespace MiniMaxH3 { ggml_tensor* time_embedding(GGMLRunnerContext* ctx, ggml_tensor* timestep_features, + ggml_tensor* endpoint_features, ggml_tensor* curve_indices, ggml_tensor* curve_upper_indices, ggml_tensor* curve_fractions) { if (!config.uses_adaln_curves()) { - return std::dynamic_pointer_cast(blocks["time_embedder"])->forward(ctx, timestep_features); + auto current = std::dynamic_pointer_cast(blocks["time_embedder"])->forward(ctx, timestep_features); + if (!config.two_time) { + return current; + } + GGML_ASSERT(endpoint_features != nullptr); + auto endpoint = std::dynamic_pointer_cast(blocks["endpoint_time_embedder"]) + ->forward(ctx, endpoint_features); + auto delta = ggml_mul(ctx->ggml_ctx, ggml_sub(ctx->ggml_ctx, endpoint, current), params["hyperflow_gate"]); + return ggml_add(ctx->ggml_ctx, current, delta); } auto lower = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_indices); auto upper = ggml_get_rows(ctx->ggml_ctx, params["adaln_t_table"], curve_upper_indices); @@ -689,6 +801,7 @@ namespace MiniMaxH3 { const std::vector& condition_audios, ggml_tensor* position_ids, ggml_tensor* timestep_features, + ggml_tensor* endpoint_features, ggml_tensor* curve_indices, ggml_tensor* curve_upper_indices, ggml_tensor* curve_fractions, @@ -784,6 +897,7 @@ namespace MiniMaxH3 { auto t_emb = time_embedding(ctx, timestep_features, + endpoint_features, curve_indices, curve_upper_indices, curve_fractions); @@ -826,7 +940,9 @@ namespace MiniMaxH3 { std::vector sequence_segments; TokenModulationSpan video_segment{}; TokenModulationSpan audio_segment{}; + // Distinct (t, r) rows of the modulation table; r == t for single-time checkpoints. std::vector timesteps; + std::vector endpoints; }; static float video_span(int64_t frame) { @@ -846,13 +962,15 @@ namespace MiniMaxH3 { return result; } - static int find_or_add_timestep(std::vector* values, float value) { - auto it = std::find(values->begin(), values->end(), value); - if (it != values->end()) { - return static_cast(it - values->begin()); + static int find_or_add_timestep(PackedSequenceLayout* layout, float value, float endpoint) { + for (size_t i = 0; i < layout->timesteps.size(); ++i) { + if (layout->timesteps[i] == value && layout->endpoints[i] == endpoint) { + return static_cast(i); + } } - values->push_back(value); - return static_cast(values->size() - 1); + layout->timesteps.push_back(value); + layout->endpoints.push_back(endpoint); + return static_cast(layout->timesteps.size() - 1); } static PackedSequenceLayout build_layout(int64_t text_len, @@ -866,7 +984,9 @@ namespace MiniMaxH3 { const std::vector& reference_blocks, const sd::Tensor& text_tags, float video_t, - float audio_timestep) { + float audio_timestep, + float video_r, + float audio_r) { PackedSequenceLayout layout; float sqrt_area = std::sqrt(static_cast(latent_h * latent_w)); auto h_axis = spatial_axis(latent_h, sqrt_area); @@ -884,12 +1004,13 @@ namespace MiniMaxH3 { } layout.sequence_segments.push_back({0, text_len, SequenceKind::TEXT}); - int video_time_row = find_or_add_timestep(&layout.timesteps, video_t); - int audio_time_row = find_or_add_timestep(&layout.timesteps, audio_timestep); - int condition_time_row = find_or_add_timestep(&layout.timesteps, - std::max(video_t, VISUAL_COND_TIMESTEP)); - int audio_condition_time_row = find_or_add_timestep(&layout.timesteps, - std::max(audio_timestep, 1.f)); + // Conditions stay at their own noise level, so their endpoint equals their timestep. + float condition_t = std::max(video_t, VISUAL_COND_TIMESTEP); + float audio_condition_t = std::max(audio_timestep, 1.f); + int video_time_row = find_or_add_timestep(&layout, video_t, video_r); + int audio_time_row = find_or_add_timestep(&layout, audio_timestep, audio_r); + int condition_time_row = find_or_add_timestep(&layout, condition_t, condition_t); + int audio_condition_time_row = find_or_add_timestep(&layout, audio_condition_t, audio_condition_t); int64_t run_start = 0; int current_tag = text_tags.empty() ? 1 : text_tags[0]; @@ -1056,6 +1177,7 @@ namespace MiniMaxH3 { MiniMaxH3Transformer3DModel model; sd::Tensor position_input_cache; sd::Tensor timestep_feature_input_cache; + sd::Tensor endpoint_feature_input_cache; sd::Tensor curve_index_input_cache; sd::Tensor curve_upper_index_input_cache; sd::Tensor curve_fraction_input_cache; @@ -1128,7 +1250,8 @@ namespace MiniMaxH3 { const std::vector& reference_blocks, int audio_length, float video_shift, - float audio_shift) { + float audio_shift, + float video_sigma_next) { GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); GGML_ASSERT(audio_length > 0); GGML_ASSERT(!context_tensor.empty()); @@ -1166,6 +1289,17 @@ namespace MiniMaxH3 { float audio_scale = video_shift / audio_shift; float t_v = 1.f - sigma_v; float t_a = 1.f - sigma_a; + float r_v = t_v; + float r_a = t_a; + if (config.two_time) { + if (std::isfinite(video_sigma_next)) { + float next = std::clamp(video_sigma_next, 0.f, 1.f); + r_v = 1.f - next; + r_a = 1.f - time_shift_sigma(next, video_shift, audio_shift); + } else { + LOG_WARN("two-time checkpoint without a step endpoint; using r = t"); + } + } // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed // 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); @@ -1180,7 +1314,9 @@ namespace MiniMaxH3 { reference_blocks, text_tags, t_v, - t_a); + t_a, + r_v, + r_a); position_input_cache = sd::Tensor( {3, static_cast(layout.positions.size() / 3)}, @@ -1188,6 +1324,7 @@ namespace MiniMaxH3 { auto positions = make_input(position_input_cache); ggml_tensor* timestep_features = nullptr; + ggml_tensor* endpoint_features = nullptr; ggml_tensor* curve_indices = nullptr; ggml_tensor* curve_upper_indices = nullptr; ggml_tensor* curve_fractions = nullptr; @@ -1224,6 +1361,16 @@ namespace MiniMaxH3 { true, 1.f)); timestep_features = make_input(timestep_feature_input_cache); + if (config.two_time) { + endpoint_feature_input_cache = sd::Tensor( + {config.timestep_input_dim, static_cast(layout.endpoints.size())}, + timestep_embedding(layout.endpoints, + static_cast(config.timestep_input_dim), + 10000, + true, + 1.f)); + endpoint_features = make_input(endpoint_feature_input_cache); + } } auto graph = new_graph_custom(H3_GRAPH_SIZE); @@ -1236,6 +1383,7 @@ namespace MiniMaxH3 { audio_condition_inputs, positions, timestep_features, + endpoint_features, curve_indices, curve_upper_indices, curve_fractions, @@ -1279,7 +1427,8 @@ namespace MiniMaxH3 { reference_blocks, extra->audio_length, extra->video_sigma_shift, - extra->audio_sigma_shift); + extra->audio_sigma_shift, + extra->video_sigma_next); }; return restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, diff --git a/src/model/diffusion/model.hpp b/src/model/diffusion/model.hpp index f25ca7b519..ed20c80870 100644 --- a/src/model/diffusion/model.hpp +++ b/src/model/diffusion/model.hpp @@ -145,6 +145,7 @@ struct MiniMaxH3DiffusionExtra { int audio_length = 0; float video_sigma_shift = 12.f; float audio_sigma_shift = 3.f; + float video_sigma_next = NAN; }; struct MiniT2IDiffusionExtra { diff --git a/src/model_manager_prefetch.cpp b/src/model_manager_prefetch.cpp index 128cae237a..4056b8f0ea 100644 --- a/src/model_manager_prefetch.cpp +++ b/src/model_manager_prefetch.cpp @@ -33,10 +33,7 @@ void ModelManager::synchronize_prefetch_block(PrefetchBlock& block) { ggml_backend_event_synchronize(block.event); ggml_backend_event_free(block.event); block.event = nullptr; - } - // Event completion does not release backend-owned transfer resources, such - // as Metal's temporary host mappings and command buffers. - if (block.submitted && block.transfer_backend != nullptr) { + } else if (block.submitted && block.transfer_backend != nullptr) { ggml_backend_synchronize(block.transfer_backend); } block.transfer_backend = nullptr; diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index c820a8d963..2a6683f010 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2593,14 +2593,16 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr(std::abs(step)); + diffusion_params.extra = MiniMaxH3DiffusionExtra{ condition.c_token_types.empty() ? nullptr : &condition.c_token_types, condition.c_position_ids.empty() ? nullptr : &condition.c_position_ids, condition.c_ref_audios.empty() ? nullptr : &condition.c_ref_audios, condition.c_reference_blocks.empty() ? nullptr : &condition.c_reference_blocks, version == VERSION_P_H3_IMAGE ? 2 : audio_length, std::isfinite(active_flow_shift) ? active_flow_shift : 12.f, - version == VERSION_P_H3_IMAGE ? 1.f : 3.f}; + version == VERSION_P_H3_IMAGE ? 1.f : 3.f, + target_index > 0 && target_index < sigmas.size() ? sigmas[target_index] : NAN}; } else if (sd_version_is_ltxav(version)) { diffusion_params.extra = LTXAVDiffusionExtra{ nullptr, From abc6718455ebf7a0494d8890a159748078f65f47 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 01:10:48 +0800 Subject: [PATCH 34/57] fix(metal): reclaim completed transfer resources without blocking --- ggml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ggml b/ggml index 1e5639ec8f..d803c238f6 160000 --- a/ggml +++ b/ggml @@ -1 +1 @@ -Subproject commit 1e5639ec8f3aa05c17ecddb039b86a09134ff03c +Subproject commit d803c238f67915fb825a9aa4df27b99356d51597 From cd5594e9f9e5964388b716036b319f1a81dd7c7e Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 01:13:57 +0800 Subject: [PATCH 35/57] feat: add SeedVR2 image and video upscaling (#2300) --- CONTRIBUTING.md | 12 ++ README.md | 4 +- docs/relay_vsr.md | 88 ++++----- docs/seedvr2.md | 69 +++++++ examples/cli/main.cpp | 72 ++++--- examples/common/common.cpp | 2 +- include/stable-diffusion.h | 5 + src/model/diffusion/seedvr2.hpp | 325 ++++++++++++++++++++++++++++++++ src/model/vae/seedvr2_vae.hpp | 286 ++++++++++++++++++++++++++++ src/name_conversion.cpp | 17 +- src/pipeline/relay_vsr.cpp | 2 +- src/pipeline/seedvr2.cpp | 167 ++++++++++++++++ src/pipeline/video_upscale.cpp | 67 ++++++- src/pipeline/video_upscale.h | 14 ++ 14 files changed, 1038 insertions(+), 92 deletions(-) create mode 100644 docs/seedvr2.md create mode 100644 src/model/diffusion/seedvr2.hpp create mode 100644 src/model/vae/seedvr2_vae.hpp create mode 100644 src/pipeline/seedvr2.cpp diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index ca64ad1981..15045dcc7e 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -29,6 +29,18 @@ PRs should include: - Linked issue/PR context when applicable. - Screenshots or sample outputs for UI/visual behavior changes. +## Documentation + +Write model usage guides in `docs/` for users. Focus on supported weights, +setup, runnable commands, parameter meanings, output behavior, and practical +limits such as memory requirements. Include technical details only when they +help users choose settings or troubleshoot a problem. + +Keep tensor layouts, internal class names, kernel implementation details, +algorithm derivations, and development history out of usage guides. Keep task-specific +profiling logs, numerical comparison procedures, and test artifacts local, and +summarize relevant verification results in the PR description. + ## README UI Listings The [UIs list in the README](README.md#uis) only includes open-source projects that use `stable-diffusion.cpp` as a backend. Do not submit PRs to add closed-source UIs. diff --git a/README.md b/README.md index 2442865b2f..3a65ee1bfe 100644 --- a/README.md +++ b/README.md @@ -88,6 +88,9 @@ API and command-line option may change frequently.*** - [LTX-2.3/LTX-2.5](./docs/ltx2.md) - [HunyuanVideo 1.5](./docs/hunyuan_video.md) - [LingBot-Video](./docs/lingbot_video.md) + - Video Upscale Models + - [RelayVSR](./docs/relay_vsr.md) + - [SeedVR2](./docs/seedvr2.md) - [PhotoMaker](./docs/photo_maker.md) support. - [IP-Adapter](./docs/ip_adapter.md) support (SD 1.5 and SDXL, including Plus) - Control Net support with SD 1.5 @@ -96,7 +99,6 @@ API and command-line option may change frequently.*** - Latent Consistency Models support (LCM/LCM-LoRA) - Faster and memory efficient latent decoding with [TAESD](./docs/taesd.md) - Upscale images generated with [ESRGAN](./docs/esrgan.md) - - Upscale videos with [RelayVSR](./docs/relay_vsr.md) (experimental, 2x and 4x) - Supported backends - CPU (AVX, AVX2 and AVX512 support for x86 architectures) - CUDA diff --git a/docs/relay_vsr.md b/docs/relay_vsr.md index 1352f41f19..e35f5fdb7f 100644 --- a/docs/relay_vsr.md +++ b/docs/relay_vsr.md @@ -1,24 +1,24 @@ -# RelayVSR video super-resolution +# RelayVSR image and video super-resolution -RelayVSR uses a Wan model to sample sparse reference latents, followed by the -Flash conditional decoder to reconstruct every input frame. The model code -follows [RelayVSR](https://github.com/kopperx/RelayVSR); the loader maps -[ComfyUI repackaged weights](https://huggingface.co/Comfy-Org/RelayVSR) to the -original module names. +[RelayVSR](https://github.com/kopperx/RelayVSR) upscales images and videos by +2x or 4x with `-M vid_upscale`. -Use `-M vid_upscale` with an ordered directory of RGB images. Filenames are -sorted lexicographically. Supply the source frame rate with `--fps`. Output -keeps the input frame count. Use a video output supported by the CLI, such as -`.avi` or `.webm`, or a numbered PNG output such as `output/%04d.png`. +## Weights -The scale is an explicit, required argument. It is not read from checkpoint -metadata. Use matching Wan and Flash checkpoints: +[ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/RelayVSR) load +directly. Select the matching pair for the requested scale: -| Scale | `--diffusion-model` | `--vae` | +| `--upscale-factor` | `--diffusion-model` | `--vae` | | --- | --- | --- | | 2 | `relayvsr_x2_lightx2v_int8_convrot.safetensors` | `relayvsr_flash_vae_x2_fp32.safetensors` | | 4 | `relayvsr_x4_int8_convrot.safetensors` | `relayvsr_flash_vae_x4_fp32.safetensors` | +No separate text encoder or ordinary Wan VAE is needed. + +## Usage + +For a video, extract numbered frames into a directory, then run: + ```sh sd-cli -M vid_upscale \ --diffusion-model models/diffusion_models/relayvsr_x2_lightx2v_int8_convrot.safetensors \ @@ -27,44 +27,24 @@ sd-cli -M vid_upscale \ --vsr-window-size 16 -s 42 -o output.avi ``` -`--vsr-window-size` includes both reference endpoints. The default 16 samples -references at frame indices 0, 15, 30, and so on, always including the last -frame. Shared endpoints update Flash history in the following window but are -emitted only once. Single-frame input is also supported. - -Wan performs one flow-matching Euler step at timestep 1000 with guidance 1, -using the prompt embedding stored in the checkpoint. No text encoder or -ordinary Wan VAE is needed. Generic generation sampling settings do not -configure this fixed sampling path. - -The Flash scale cannot be verified from its tensor shapes: the x2 and x4 files -have identical shapes. The caller must select the correct pair. Wan's input -projection is checked against the requested scale. Inputs are replicate-padded -internally and outputs are cropped to exactly the input size times the scale. - -INT8 convrot uses the existing [quantized-weight execution path](int8_convrot.md). -Backend and parameter placement use `--backend` and `--params-backend`, with -Wan assigned to `diffusion` and Flash to `vae`. CUDA builds with CUTLASS use -fused window attention on compute capability 8.0 or newer. It loads spatial -neighborhoods directly, with FP16 Q/K/V and FP32 accumulation, and preserves -Wan's boundary-shifted windows and Flash's boundary-truncated windows. -Other backends use bounded spatial tiles with ordinary attention operators. -Set `GGML_CUDA_DISABLE_WINDOW_ATTENTION=1` to use that fallback on CUDA for -comparison. Floating-point results can differ between these implementations. - -Library callers initialize `sd_video_upscaler_ctx_params_t` with -`sd_video_upscaler_ctx_params_init()`, set the model paths, explicit scale and -backend options, and create a `video_upscaler_ctx_t` with -`new_video_upscaler_ctx()`. Model architecture is detected from -weight names and shapes; the scale remains caller-supplied. - -Initialize `sd_video_upscale_params_t` with `sd_video_upscale_params_init()`, -then set the RGB frames, frame count, FPS, seed and RelayVSR reference window. -Call `upscale_video()` and free the result with -`free_sd_video_generation_result()`. Defaults are 24 FPS, seed 42 and a -16-frame reference window. No sampler or scheduler selection is required. - -Each call starts fresh -temporal state. Model caches are bounded by the temporal windows; the public -API currently holds the complete input and output frame arrays in host memory. -Audio is not copied by this path. +Frames are read in filename order. Set `--fps` to the source video's frame rate. +Video outputs support AVI, WebM, and animated WebP. To save individual frames, +use an output pattern such as `-o output/%04d.png`. + +For a single image, replace `--input-frames input_frames` with `-i input.png` +and use `-o output.png` (or JPEG/WebP). No FPS setting is needed for an image. +`-i` / `--init-img` and `--input-frames` are mutually exclusive. + +## Settings and limits + +- `--upscale-factor` is required. Use 2 or 4 with the corresponding weights. + Output width and height are exactly the input dimensions times the scale. +- `--vsr-window-size` sets the reference-frame window, including both endpoints. + The default 16 selects reference frames 0, 15, 30, and so on, and always + includes the final frame. Every input frame appears once in the output. +- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See + [backend selection](backend.md) to choose a device or weight placement. +- Sampler, scheduler, and step settings do not configure RelayVSR. + `--vsr-temporal-chunk-size` is a SeedVR2 option. +- Long videos require enough system memory to hold the input and output frames. + Audio is not copied to the output. diff --git a/docs/seedvr2.md b/docs/seedvr2.md new file mode 100644 index 0000000000..24f55032d6 --- /dev/null +++ b/docs/seedvr2.md @@ -0,0 +1,69 @@ +# SeedVR2 image and video super-resolution + +SeedVR2 upscales images and videos with `-M vid_upscale`. Support is experimental. + +## Weights + +These [ComfyUI safetensors weights](https://huggingface.co/Comfy-Org/SeedVR2) +load directly, without conversion or a separate text encoder: + +| Model | `--diffusion-model` | `--vae` | +| --- | --- | --- | +| 3B | `seedvr2_3b_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` | +| 7B Sharp | `seedvr2_7b_sharp_int8_convrot.safetensors` | `seedvr2_ema_vae_fp16.safetensors` | + +Use the matching SeedVR2 VAE for either model. The 7B model requires more memory. + +## Upscale an image + +```sh +sd-cli -M vid_upscale \ + --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \ + --vae models/vae/seedvr2_ema_vae_fp16.safetensors \ + -i input.png --upscale-factor 2 -s 42 -o restored.png +``` + +`-i` / `--init-img` accepts a single image. Outputs can be PNG, JPEG, or WebP. +No FPS or temporal chunk setting is needed for an image. + +## Upscale a video + +Extract the video into numbered frames, then pass the frame directory: + +```sh +mkdir frames +ffmpeg -i input.mp4 -fps_mode passthrough frames/%06d.png +sd-cli -M vid_upscale \ + --diffusion-model models/diffusion_models/seedvr2_3b_int8_convrot.safetensors \ + --vae models/vae/seedvr2_ema_vae_fp16.safetensors \ + --input-frames frames --upscale-factor 2 --fps 24 \ + --vsr-temporal-chunk-size 1 -s 42 -o restored.avi +``` + +Set `--fps` to the source video's frame rate; 24 is only the example above. +Frames are read in filename order. Use either `-i` or `--input-frames`, not both. + +Video outputs support AVI, WebM, and animated WebP. To save individual frames, +use an output pattern such as `-o restored/%06d.png`. + +## Settings and limits + +- `--upscale-factor` is required and must be a positive integer. Output width + and height are exactly the input dimensions multiplied by this value. +- `--vsr-temporal-chunk-size` controls how many video frames are processed + together. Valid values are 1, 5, 9, 13, and so on; the default is 5. A value of + 1 processes each frame independently and uses less memory. Larger values use + neighboring frames but require more memory. Chunk boundaries may be visible. +- Frame count and order are preserved, including a shorter final chunk. Audio + is not copied to the output. +- If you run out of GPU memory, try `--vsr-temporal-chunk-size 1`. A five-frame + chunk at 1728x960 output resolution exceeded 24 GB of GPU memory in testing. +- `--flash-attn auto` and `--conv-direct auto` are enabled by default. See + [backend selection](backend.md) to choose a device or weight placement. +- SeedVR2 does not use sampler, scheduler, or step settings. + `--vsr-window-size` is a RelayVSR option and does not affect SeedVR2. + +3B has been tested on CUDA and CPU; 7B has been tested on CUDA with small inputs. +Other backends have not been tested with this integration. + +Model reference: [ByteDance-Seed/SeedVR](https://github.com/ByteDance-Seed/SeedVR). diff --git a/examples/cli/main.cpp b/examples/cli/main.cpp index bcf73595c6..2bff0507f7 100644 --- a/examples/cli/main.cpp +++ b/examples/cli/main.cpp @@ -35,12 +35,13 @@ std::regex format_specifier_regex("(?:[^%]|^)(?:%%)*(%\\d{0,3}d)"); struct SDCliParams { std::string input_frames; - int video_upscale_factor = 0; - int vsr_window_size = 16; - SDMode mode = IMG_GEN; - std::string output_path = "output.png"; - int output_begin_idx = -1; - int compression_quality = 90; + int video_upscale_factor = 0; + int vsr_window_size = 16; + int vsr_temporal_chunk_size = 5; + SDMode mode = IMG_GEN; + std::string output_path = "output.png"; + int output_begin_idx = -1; + int compression_quality = 90; std::string image_path; std::string metadata_format = "text"; @@ -66,7 +67,7 @@ struct SDCliParams { ArgOptions options; options.string_options = { - {"", "--input-frames", "directory of ordered RGB frames for vid_upscale", 0, &input_frames}, + {"", "--input-frames", "directory of ordered RGB frames for vid_upscale (use -i for a single image)", 0, &input_frames}, {"-o", "--output", "path to write result image to. you can use printf-style %d format specifiers for image sequences (default: ./output.png) (eg. output_%03d.png). Single-file video outputs support .avi, .webm, and animated .webp", @@ -95,8 +96,9 @@ struct SDCliParams { }; options.int_options = { - {"", "--upscale-factor", "required video upscale factor (RelayVSR: 2 or 4)", &video_upscale_factor}, + {"", "--upscale-factor", "required positive video upscale factor (RelayVSR: 2 or 4; SeedVR2: integer)", &video_upscale_factor}, {"", "--vsr-window-size", "RelayVSR window including both endpoints (default: 16)", &vsr_window_size}, + {"", "--vsr-temporal-chunk-size", "SeedVR2 input frames per independent chunk: 1 or 4n+1 (default: 5)", &vsr_temporal_chunk_size}, {"", "--preview-interval", "preview interval: in each sampling pass, positive N updates every Nth denoiser step and -N previews only completed logical step N; 0 previews the final completed step of the first pass (base-resolution or high-noise). Default: 1", @@ -227,8 +229,8 @@ struct SDCliParams { } bool validate() { - if (mode == VID_UPSCALE && (input_frames.empty() || video_upscale_factor <= 0 || vsr_window_size < 2)) { - LOG_ERROR("vid_upscale requires --input-frames, positive --upscale-factor, and --vsr-window-size >= 2"); + if (mode == VID_UPSCALE && video_upscale_factor <= 0) { + LOG_ERROR("vid_upscale requires a positive --upscale-factor"); return false; } if (mode != METADATA) { @@ -540,7 +542,9 @@ bool save_results(const SDCliParams& cli_params, return sucessful_reults != 0; } - if ((cli_params.mode == VID_GEN && num_results > 1) || cli_params.mode == VID_UPSCALE) { + const bool video_upscale_output = cli_params.mode == VID_UPSCALE && + (gen_params.init_image_path.empty() || ext_lower == ".avi" || ext_lower == ".webm"); + if ((cli_params.mode == VID_GEN && num_results > 1) || video_upscale_output) { if (ext_lower != ".avi" && ext_lower != ".webp" && ext_lower != ".webm") ext = ".avi"; std::string params = gen_params.embed_image_metadata @@ -676,12 +680,23 @@ int main(int argc, const char* argv[]) { } if (cli_params.mode == VID_UPSCALE) { + if (cli_params.input_frames.empty() == gen_params.init_image_path.empty()) { + LOG_ERROR("vid_upscale requires exactly one input: -i/--init-img or --input-frames"); + return 1; + } + ctx_params.prepare_backend_assignments(); if (ctx_params.diffusion_model_path.empty() || ctx_params.vae_path.empty() || gen_params.fps <= 0) { LOG_ERROR("vid_upscale requires --diffusion-model, --vae and positive --fps"); return 1; } std::vector owned_frames; - if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) { + if (!gen_params.init_image_path.empty()) { + owned_frames.emplace_back(); + if (!load_sd_image_from_file(owned_frames.back().put(), gen_params.init_image_path.c_str(), 0, 0, 3)) { + LOG_ERROR("load image from '%s' failed", gen_params.init_image_path.c_str()); + return 1; + } + } else if (!load_images_from_dir(cli_params.input_frames, owned_frames) || owned_frames.empty()) { LOG_ERROR("vid_upscale input contains no readable frames"); return 1; } @@ -692,12 +707,22 @@ int main(int argc, const char* argv[]) { frames.push_back(frame.get()); sd_video_upscaler_ctx_params_t params; sd_video_upscaler_ctx_params_init(¶ms); - params.diffusion_model_path = ctx_params.diffusion_model_path.c_str(); - params.vae_path = ctx_params.vae_path.c_str(); - params.upscale_factor = cli_params.video_upscale_factor; - params.n_threads = ctx_params.n_threads; - params.backend = ctx_params.effective_backend.c_str(); - params.params_backend = ctx_params.effective_params_backend.c_str(); + params.diffusion_model_path = ctx_params.diffusion_model_path.c_str(); + params.vae_path = ctx_params.vae_path.c_str(); + params.upscale_factor = cli_params.video_upscale_factor; + params.n_threads = ctx_params.n_threads; + params.backend = ctx_params.effective_backend.c_str(); + params.params_backend = ctx_params.effective_params_backend.c_str(); + params.flash_attn_mode = ctx_params.flash_attn_mode.c_str(); + std::string conv_direct_mode = ctx_params.conv_direct_mode; + if (conv_direct_mode.empty()) { + conv_direct_mode = "auto"; + if (ctx_params.diffusion_conv_direct) + conv_direct_mode += ",diffusion=on"; + if (ctx_params.vae_conv_direct) + conv_direct_mode += ",vae=on"; + } + params.conv_direct_mode = conv_direct_mode.c_str(); std::unique_ptr upscaler( new_video_upscaler_ctx(¶ms), free_video_upscaler_ctx); @@ -706,11 +731,12 @@ int main(int argc, const char* argv[]) { sd_video_generation_result_t* output = nullptr; sd_video_upscale_params_t upscale_params; sd_video_upscale_params_init(&upscale_params); - upscale_params.frames = frames.data(); - upscale_params.frame_count = int(frames.size()); - upscale_params.fps = gen_params.fps; - upscale_params.window_size = cli_params.vsr_window_size; - upscale_params.seed = gen_params.seed; + upscale_params.frames = frames.data(); + upscale_params.frame_count = int(frames.size()); + upscale_params.fps = gen_params.fps; + upscale_params.window_size = cli_params.vsr_window_size; + upscale_params.temporal_chunk_size = cli_params.vsr_temporal_chunk_size; + upscale_params.seed = gen_params.seed; if (!upscale_video(upscaler.get(), &upscale_params, &output)) return 1; std::string video_info = output->generation_info; diff --git a/examples/common/common.cpp b/examples/common/common.cpp index a3496f4fcc..70107a396f 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1136,7 +1136,7 @@ ArgOptions SDGenerationParams::get_options() { &extra_ad_args}, {"-i", "--init-img", - "path to the init image", + "path to the init image, or a single input image for vid_upscale", 0, &init_image_path}, {"", diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index a1d82e783a..53c6b7059c 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -625,6 +625,9 @@ typedef struct { int n_threads; const char* backend; const char* params_backend; + // auto/on/off or diffusion/vae assignments; null selects auto. + const char* flash_attn_mode; + const char* conv_direct_mode; } sd_video_upscaler_ctx_params_t; typedef struct { @@ -634,6 +637,8 @@ typedef struct { int64_t seed; // RelayVSR reference window, including both endpoints; defaults to 16. int window_size; + // SeedVR2 independent temporal chunks in input frames (1 or 4n+1); defaults to 5. + int temporal_chunk_size; } sd_video_upscale_params_t; SD_API void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params); diff --git a/src/model/diffusion/seedvr2.hpp b/src/model/diffusion/seedvr2.hpp new file mode 100644 index 0000000000..87e542e910 --- /dev/null +++ b/src/model/diffusion/seedvr2.hpp @@ -0,0 +1,325 @@ +#ifndef __SD_MODEL_DIFFUSION_SEEDVR2_HPP__ +#define __SD_MODEL_DIFFUSION_SEEDVR2_HPP__ + +#include "model/common/vsr.hpp" + +namespace SeedVR2 { + + struct SeedVR2Config { + int dim = 0; + int layers = 0; + int mm_layers = 0; + int hidden = 0; + int text_dim = 0; + int text_len = 0; + int rope_freqs = 0; + bool swiglu = false; + + static SeedVR2Config detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) { + SeedVR2Config config; + auto at = [&](const std::string& key) -> const TensorStorage& { + auto found = weights.find(prefix + "." + key); + if (found == weights.end()) + throw std::runtime_error("Missing SeedVR2 tensor: " + prefix + "." + key); + return found->second; + }; + config.dim = int(at("vid_in.proj.weight").ne[1]); + config.hidden = int(at("blocks.0.mlp.vid.proj_in.weight").ne[1]); + config.text_dim = int(at("positive_conditioning").ne[0]); + config.text_len = int(at("positive_conditioning").ne[1]); + config.rope_freqs = int(at("blocks.0.attn.rope.rope.freqs").ne[0]); + config.swiglu = weights.count(prefix + ".blocks.0.mlp.vid.proj_in_gate.weight") != 0; + while (weights.count(prefix + ".blocks." + std::to_string(config.layers) + ".attn.rope.rope.freqs")) { + if (weights.count(prefix + ".blocks." + std::to_string(config.layers) + ".attn.proj_qkv.vid.weight")) + ++config.mm_layers; + ++config.layers; + } + if (at("vid_in.proj.weight").logical_ne0() != 132 || at("vid_out.proj.weight").ne[1] != 64 || + config.text_dim != 5120 || + !(config.swiglu ? config.dim == 2560 && config.layers == 32 && config.mm_layers == 10 && config.rope_freqs == 21 + : config.dim == 3072 && config.layers == 36 && config.mm_layers == 36 && config.rope_freqs == 10)) + throw std::runtime_error("Unsupported SeedVR2 architecture"); + LOG_VERBOSE("seedvr2: dim = %d, layers = %d, mm_layers = %d, hidden = %d, rope_freqs = %d", config.dim, config.layers, config.mm_layers, config.hidden, config.rope_freqs); + return config; + } + }; + + struct WindowPlan { + sd::Tensor order; + sd::Tensor reverse; + sd::Tensor positions; + sd::Tensor text_positions; + std::vector lengths; + + WindowPlan(int width, int height, int frames, int text_len, bool shifted, bool multimodal) + : order({width * height * frames}), reverse({width * height * frames}), positions({1, 3, width * height * frames}), text_positions({1, 3, text_len}) { + const double scale = std::sqrt(720.0 * 1280.0 / 256.0 / (width * height)); + const int wh = int(std::ceil(std::nearbyint(height * scale) / 3)); + const int ww = int(std::ceil(std::nearbyint(width * scale) / 3)); + const int wt = (std::min(frames, 30) + 3) / 4; + const double st = shifted && wt < frames ? 0.5 : 0; + const double sh = shifted && wh < height ? 0.5 : 0; + const double sw = shifted && ww < width ? 0.5 : 0; + const int nt = st ? int(std::ceil((frames - st) / wt)) + 1 : (frames + wt - 1) / wt; + const int nh = sh ? int(std::ceil((height - sh) / wh)) + 1 : (height + wh - 1) / wh; + const int nw = sw ? int(std::ceil((width - sw) / ww)) + 1 : (width + ww - 1) / ww; + int index = 0; + for (int iw = 0; iw < nw; ++iw) + for (int ih = 0; ih < nh; ++ih) + for (int it = 0; it < nt; ++it) { + int x0 = std::max(0, int((iw - sw) * ww)), x1 = std::min(width, int((iw - sw + 1) * ww)); + int y0 = std::max(0, int((ih - sh) * wh)), y1 = std::min(height, int((ih - sh + 1) * wh)); + int t0 = std::max(0, int((it - st) * wt)), t1 = std::min(frames, int((it - st + 1) * wt)); + if (x1 <= x0 || y1 <= y0 || t1 <= t0) + continue; + lengths.push_back((x1 - x0) * (y1 - y0) * (t1 - t0)); + for (int t = t0; t < t1; ++t) + for (int y = y0; y < y1; ++y) + for (int x = x0; x < x1; ++x, ++index) { + int original = (t * height + y) * width + x; + order.data()[index] = original; + reverse.data()[original] = index; + const int pos[] = {t - t0, y - y0, x - x0}; + const int size[] = {t1 - t0, y1 - y0, x1 - x0}; + for (int a = 0; a < 3; ++a) + positions.data()[3 * index + a] = multimodal ? float(pos[a] + (a == 0 ? text_len : 0)) + : (size[a] == 1 ? -1.f : -1.f + 2.f * pos[a] / (size[a] - 1)); + } + } + if (index != width * height * frames) + throw std::runtime_error("SeedVR2 window partition is incomplete"); + for (int i = 0; i < text_len; ++i) + for (int a = 0; a < 3; ++a) + text_positions.data()[3 * i + a] = multimodal ? float(i) : 0.f; + } + }; + + class AdaSingle : public GGMLBlock { + int dim_; + bool output_; + + protected: + void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override { + for (const auto& layer : output_ ? std::vector{"out"} : std::vector{"attn", "mlp"}) { + for (const auto& mode : output_ ? std::vector{"shift", "scale"} : std::vector{"shift", "scale", "gate"}) + params[layer + "_" + mode] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim_); + } + } + + public: + AdaSingle(int dim, bool output = false) + : dim_(dim), output_(output) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* emb, const std::string& layer, bool gate = false) { + auto c = ctx->ggml_ctx; + // The reference reshapes embedding to [D, 2, 3], not [6, D]. + const int offset = layer == "mlp" ? 3 : 0; + auto mod = [&](int i, const char* name) { + auto value = ggml_cont(c, ggml_view_2d(c, emb, 1, dim_, 6 * emb->nb[0], (offset + i) * emb->nb[0])); + return ggml_add(c, ggml_reshape_1d(c, value, dim_), params[layer + "_" + name]); + }; + if (gate) + return ggml_mul(c, x, mod(2, "gate")); + // AdaSingle uses scale directly; ggml_ext_modulate adds an identity scale. + return ggml_add(c, ggml_mul(c, x, mod(1, "scale")), mod(0, "shift")); + } + }; + + class SRMLP : public UnaryBlock { + bool swiglu_; + + public: + explicit SRMLP(const SeedVR2Config& config) + : swiglu_(config.swiglu) { + blocks["proj_in"] = std::make_shared(config.dim, config.hidden, !swiglu_); + blocks["proj_out"] = std::make_shared(config.hidden, config.dim, !swiglu_); + if (swiglu_) + blocks["proj_in_gate"] = std::make_shared(config.dim, config.hidden, false); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto up = std::dynamic_pointer_cast(blocks["proj_in"])->forward(ctx, x); + auto hidden = swiglu_ ? ggml_swiglu_split(ctx->ggml_ctx, std::dynamic_pointer_cast(blocks["proj_in_gate"])->forward(ctx, x), up) + : ggml_ext_gelu(ctx->ggml_ctx, up); + return std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, hidden); + } + }; + + class NaMMSRTransformerBlock : public GGMLBlock { + SeedVR2Config config_; + bool shared_; + bool last_; + + protected: + void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override { + params["attn.rope.rope.freqs"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, config_.rope_freqs); + } + + public: + NaMMSRTransformerBlock(const SeedVR2Config& config, int index) + : config_(config), shared_(index >= config.mm_layers), last_(config.swiglu && index == config.layers - 1) { + for (const auto& branch : shared_ ? std::vector{"all"} : std::vector{"vid", "txt"}) { + blocks["attn.proj_qkv." + branch] = std::make_shared(config.dim, 3 * config.dim, false); + blocks["attn.proj_out." + branch] = std::make_shared(config.dim, config.dim); + blocks["attn.norm_q." + branch] = std::make_shared(128, 1e-5f); + blocks["attn.norm_k." + branch] = std::make_shared(128, 1e-5f); + blocks["ada." + branch] = std::make_shared(config.dim); + blocks["mlp." + branch] = std::make_shared(config); + } + } + + ggml_tensor* rotation(GGMLRunnerContext* ctx, ggml_tensor* positions) { + auto c = ctx->ggml_ctx; + const int n = config_.rope_freqs; + auto p = ggml_repeat(c, positions, ggml_new_tensor_3d(c, GGML_TYPE_F32, n, 3, positions->ne[2])); + auto angle = ggml_mul(c, p, params["attn.rope.rope.freqs"]); + angle = ggml_reshape_2d(c, angle, 3 * n, positions->ne[2]); + angle = ggml_pad(c, angle, 64 - 3 * n, 0, 0, 0); + // Keep tokens out of ne[3]: CUDA concat otherwise launches once per token. + auto cs = ggml_reshape_3d(c, ggml_cos(c, angle), 1, 1, 64 * positions->ne[2]); + auto sn = ggml_reshape_3d(c, ggml_sin(c, angle), 1, 1, 64 * positions->ne[2]); + auto matrix = ggml_concat(c, ggml_concat(c, cs, ggml_scale(c, sn, -1.f), 0), ggml_concat(c, sn, cs, 0), 1); + return ggml_reshape_4d(c, matrix, 2, 2, 64, positions->ne[2]); + } + + std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* vid, ggml_tensor* txt, ggml_tensor* emb, const WindowPlan& plan, VSR::GraphInputs& inputs) { + auto c = ctx->ggml_ctx; + const int dim = config_.dim, heads = dim / 128; + auto branch = [&](bool text) -> std::string { return shared_ ? "all" : text ? "txt" + : "vid"; }; + auto linear = [&](const std::string& name, bool text, ggml_tensor* x) { + return std::dynamic_pointer_cast(blocks[name + "." + branch(text)])->forward(ctx, x); + }; + auto ada = [&](ggml_tensor* x, bool text, const std::string& layer, bool gate = false) { + if (text && last_) + return x; + return std::dynamic_pointer_cast(blocks["ada." + branch(text)])->forward(ctx, x, emb, layer, gate); + }; + auto vqkv = linear("attn.proj_qkv", false, ada(ggml_rms_norm(c, vid, 1e-5f), false, "attn")); + auto tqkv = linear("attn.proj_qkv", true, ada(ggml_rms_norm(c, txt, 1e-5f), true, "attn")); + vqkv = ggml_get_rows(c, vqkv, inputs.input(ctx, plan.order)); + auto vpe = rotation(ctx, inputs.input(ctx, plan.positions)); + auto tpe = rotation(ctx, inputs.input(ctx, plan.text_positions)); + auto qkv = [&](ggml_tensor* x, bool text, ggml_tensor* pe) { + auto parts = ggml_ext_chunk(c, x, 3, 0); + for (int i = 0; i < 2; ++i) { + auto h = ggml_reshape_3d(c, parts[i], 128, heads, x->ne[1]); + h = std::dynamic_pointer_cast(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h); + h = ggml_reshape_2d(c, h, dim, x->ne[1]); + parts[i] = VSR::apply_rope(ctx, h, pe, heads); + } + return parts; + }; + auto vp = qkv(vqkv, false, vpe), tp = qkv(tqkv, true, tpe); + std::vector video_windows; + ggml_tensor* tout = nullptr; + int start = 0; + for (int length : plan.lengths) { + std::vector joined; + for (int i = 0; i < 3; ++i) + joined.push_back(ggml_concat(c, ggml_ext_slice(c, vp[i], 1, start, start + length), tp[i], 1)); + auto out = ggml_ext_attention_ext(c, ctx->backend, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled); + auto vo = ggml_ext_slice(c, out, 1, 0, length); + auto to = ggml_ext_slice(c, out, 1, length, length + txt->ne[1]); + video_windows.push_back(vo); + tout = tout ? ggml_add(c, tout, to) : to; + start += length; + } + auto vout = ggml_ext_vec_concat(c, video_windows, 1); + vout = ggml_get_rows(c, vout, inputs.input(ctx, plan.reverse)); + tout = ggml_scale(c, tout, 1.f / plan.lengths.size()); + vid = ggml_add(c, vid, ada(linear("attn.proj_out", false, vout), false, "attn", true)); + txt = ggml_add(c, txt, ada(linear("attn.proj_out", true, tout), true, "attn", true)); + auto mlp = [&](ggml_tensor* x, bool text) { + auto h = ada(ggml_rms_norm(c, x, 1e-5f), text, "mlp"); + h = std::dynamic_pointer_cast(blocks["mlp." + branch(text)])->forward(ctx, h); + return ggml_add(c, x, ada(h, text, "mlp", true)); + }; + vid = mlp(vid, false); + if (!last_) + txt = mlp(txt, true); + return {vid, txt}; + } + }; + + class NaDiT : public GGMLBlock { + SeedVR2Config config_; + + protected: + void init_params(ggml_context* ctx, const String2TensorStorage&, const std::string) override { + params["positive_conditioning"] = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, config_.text_dim, config_.text_len); + } + + public: + explicit NaDiT(const SeedVR2Config& config) + : config_(config) { + blocks["vid_in.proj"] = std::make_shared(132, config.dim); + blocks["txt_in"] = std::make_shared(config.text_dim, config.dim); + blocks["emb_in.proj_in"] = std::make_shared(256, config.dim); + blocks["emb_in.proj_hid"] = std::make_shared(config.dim, config.dim); + blocks["emb_in.proj_out"] = std::make_shared(config.dim, 6 * config.dim); + for (int i = 0; i < config.layers; ++i) + blocks["blocks." + std::to_string(i)] = std::make_shared(config, i); + blocks["vid_out.proj"] = std::make_shared(config.dim, 64); + if (config.swiglu) { + blocks["vid_out_norm"] = std::make_shared(config.dim, 1e-5f); + blocks["vid_out_ada"] = std::make_shared(config.dim, true); + } + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input, ggml_tensor* timestep, const std::vector& plans, VSR::GraphInputs& inputs) { + auto linear = [&](const std::string& name, ggml_tensor* x) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, x); }; + auto vid = linear("vid_in.proj", input); + auto txt = linear("txt_in", params["positive_conditioning"]); + auto emb = linear("emb_in.proj_in", timestep); + emb = linear("emb_in.proj_hid", ggml_silu(ctx->ggml_ctx, emb)); + emb = linear("emb_in.proj_out", ggml_silu(ctx->ggml_ctx, emb)); + for (int i = 0; i < config_.layers; ++i) { + auto name = "blocks." + std::to_string(i); + std::tie(vid, txt) = std::dynamic_pointer_cast(blocks[name])->forward(ctx, vid, txt, emb, plans[i % 2], inputs); + sd::ggml_graph_cut::mark_graph_cut(vid, "seedvr2." + name, "vid"); + sd::ggml_graph_cut::mark_graph_cut(txt, "seedvr2." + name, "txt"); + } + if (config_.swiglu) { + vid = std::dynamic_pointer_cast(blocks["vid_out_norm"])->forward(ctx, vid); + vid = std::dynamic_pointer_cast(blocks["vid_out_ada"])->forward(ctx, vid, emb, "out"); + } + return linear("vid_out.proj", vid); + } + }; + + struct SeedVR2Runner : GGMLRunner { + SeedVR2Config config; + NaDiT model; + sd::Tensor timestep; + + SeedVR2Runner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr manager) + : GGMLRunner(backend, manager), config(SeedVR2Config::detect_from_weights(weights, "diffusion")), model(config), timestep({256}) { + model.init(params_ctx, weights, "diffusion"); + for (int i = 0; i < 128; ++i) { + const double angle = 1000.0 * std::pow(10000.0, -double(i) / 128); + timestep.data()[i] = float(std::sin(angle)); + timestep.data()[i + 128] = float(std::cos(angle)); + } + } + std::string get_desc() override { return "SeedVR2 DiT"; } + void get_param_tensors(std::map& tensors) { model.get_param_tensors(tensors, "diffusion"); } + + sd::Tensor compute(int n_threads, const sd::Tensor& input, int width, int height, int frames) { + std::vector plans; + for (bool shifted : {false, true}) + plans.emplace_back(width, height, frames, config.text_len, shifted, config.swiglu); + VSR::GraphInputs inputs; + auto graph = [&]() { + inputs.begin_graph(); + auto gf = new_graph_custom(MAX_GRAPH_SIZE); + auto ctx = get_context(gf); + auto out = model.forward(&ctx, make_input(input), make_input(timestep), plans, inputs); + ggml_build_forward_expand(gf, out); + return gf; + }; + auto output = GGMLRunner::compute(graph, n_threads); + return output ? *output : sd::Tensor(); + } + }; +} + +#endif // __SD_MODEL_DIFFUSION_SEEDVR2_HPP__ diff --git a/src/model/vae/seedvr2_vae.hpp b/src/model/vae/seedvr2_vae.hpp new file mode 100644 index 0000000000..5f38aebc18 --- /dev/null +++ b/src/model/vae/seedvr2_vae.hpp @@ -0,0 +1,286 @@ +#ifndef __SD_MODEL_VAE_SEEDVR2_VAE_HPP__ +#define __SD_MODEL_VAE_SEEDVR2_VAE_HPP__ + +#include "model/common/ggml_block.hpp" + +namespace SeedVR2 { + + struct SeedVR2VAEConfig { + std::vector channels; + int latent_channels = 16; + int layers = 2; + + static SeedVR2VAEConfig detect_from_weights(const String2TensorStorage& weights, const std::string& prefix) { + SeedVR2VAEConfig config; + auto at = [&](const std::string& key) -> const TensorStorage& { + auto found = weights.find(prefix + "." + key); + if (found == weights.end()) + throw std::runtime_error("Missing SeedVR2 VAE tensor: " + prefix + "." + key); + return found->second; + }; + for (int i = 0; i < 4; ++i) + config.channels.push_back(int(at("encoder.down_blocks." + std::to_string(i) + ".resnets.0.norm2.weight").ne[0])); + if (config.channels != std::vector({128, 256, 512, 512}) || at("encoder.conv_out.bias").ne[0] != 32 || + at("decoder.conv_in.weight").ne[3] != 16 * 512 || at("decoder.conv_out.bias").ne[0] != 3 || + !weights.count(prefix + ".decoder.up_blocks.0.upsamplers.0.upscale_conv.weight")) + throw std::runtime_error("SeedVR2 requires its matching 16-channel causal VAE"); + LOG_VERBOSE("seedvr2 vae: latent_channels = %d, layers = %d, spatial factor = 8, temporal factor = 4", config.latent_channels, config.layers); + return config; + } + }; + + class InflatedCausalConv3d : public Conv3d { + int temporal_padding_; + + public: + InflatedCausalConv3d(int in, int out, std::tuple kernel = {3, 3, 3}, std::tuple stride = {1, 1, 1}, std::tuple padding = {1, 1, 1}) + : Conv3d(in, out, kernel, stride, {0, std::get<1>(padding), std::get<2>(padding)}), temporal_padding_(2 * std::get<0>(padding)) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto c = ctx->ggml_ctx; + const int kt = std::get<0>(kernel_size), kh = std::get<1>(kernel_size), kw = std::get<2>(kernel_size); + if (ctx->conv3d_direct_enabled && !ctx->weight_adapter && x->ne[2] == 1 && kt > 1 && + temporal_padding_ == kt - 1 && std::get<0>(dilation) == 1) { + // All causal taps read the same frame. Sum in F32 before rounding the + // equivalent spatial filter to F16, instead of copying and convolving the halo. + auto w = ggml_cast(c, params["weight"], GGML_TYPE_F32); + w = ggml_reshape_3d(c, w, kw * kh, kt, in_channels * out_channels); + auto spatial = ggml_ext_slice(c, w, 1, 0, 1, false); + for (int i = 1; i < kt; ++i) + spatial = ggml_add(c, spatial, ggml_ext_slice(c, w, 1, i, i + 1, false)); + w = ggml_cast(c, ggml_reshape_4d(c, spatial, kw, kh, 1, in_channels * out_channels), GGML_TYPE_F16); + if (ctx->conv2d_direct_enabled) { + w = ggml_reshape_4d(c, w, kw, kh, in_channels, out_channels); + x = ggml_reshape_4d(c, ggml_ext_cont(c, x), x->ne[0], x->ne[1], in_channels, 1); + x = ggml_ext_conv_2d(c, x, w, bias ? params["bias"] : nullptr, + std::get<2>(stride), std::get<1>(stride), + std::get<2>(padding), std::get<1>(padding), + std::get<2>(dilation), std::get<1>(dilation), true); + return ggml_reshape_4d(c, x, x->ne[0], x->ne[1], 1, out_channels); + } + return ggml_ext_conv_3d(c, ctx->backend, x, w, bias ? params["bias"] : nullptr, in_channels, + std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), + std::get<2>(padding), std::get<1>(padding), 0, + std::get<2>(dilation), std::get<1>(dilation), 1, force_prec_f32, true); + } + if (temporal_padding_) { + const int64_t w = x->ne[0], h = x->ne[1], t = x->ne[2], ch = x->ne[3]; + // Merge spatial axes so concat does not launch separately for every channel. + x = ggml_reshape_3d(c, ggml_ext_cont(c, x), w * h, t, ch); + auto head = ggml_ext_slice(c, x, 1, 0, 1); + auto shape = ggml_new_tensor_3d(c, x->type, w * h, temporal_padding_, ch); + x = ggml_concat(c, ggml_repeat(c, head, shape), x, 1); + x = ggml_reshape_4d(c, x, w, h, t + temporal_padding_, ch); + } + if (ctx->conv3d_direct_enabled) + return Conv3d::forward(ctx, x); + const int st = std::get<0>(stride), sh = std::get<1>(stride), sw = std::get<2>(stride); + const int ph = std::get<1>(padding), pw = std::get<2>(padding); + const int64_t oh = (x->ne[1] + 2 * ph - kh) / sh + 1; + const int64_t ow = (x->ne[0] + 2 * pw - kw) / sw + 1; + const int64_t ot = (x->ne[2] - kt) / st + 1; + const int64_t row_bytes = ow * in_channels * kt * kh * kw * 2; + const int64_t max_patch_bytes = 256LL * 1024 * 1024; + if (row_bytes * oh * ot <= max_patch_bytes) + return Conv3d::forward(ctx, x); + // Bound im2col storage without spatial tiling: every strip keeps the exact + // convolution halo, and normalization still sees the complete frame. + const int64_t rows = std::max(1, max_patch_bytes / row_bytes); + struct RestorePadding { + std::tuple& target; + std::tuple value; + ~RestorePadding() { target = value; } + } restore{padding, padding}; + padding = {0, 0, pw}; + auto output = ggml_ext_zeros(c, ow, oh, ot, out_channels); + for (int64_t t = 0; t < ot; ++t) { + auto frame = ggml_ext_slice(c, x, 2, t * st, t * st + kt, false); + for (int64_t y = 0; y < oh; y += rows) { + const int64_t end = std::min(oh, y + rows); + const int64_t top = y * sh - ph, bottom = (end - 1) * sh - ph + kh; + auto strip = ggml_ext_slice(c, frame, 1, std::max(0, top), std::min(x->ne[1], bottom), false); + strip = ggml_ext_pad_ext(c, ctx->backend, strip, 0, 0, int(std::max(0, -top)), + int(std::max(0, bottom - x->ne[1])), 0, 0, 0, 0); + strip = ggml_ext_cont(c, strip); + strip = Conv3d::forward(ctx, strip); + output = ggml_set_inplace(c, output, strip, output->nb[1], output->nb[2], output->nb[3], y * output->nb[1] + t * output->nb[2]); + } + } + return output; + } + }; + + static ggml_tensor* causal_norm(GGMLRunnerContext* ctx, ggml_tensor* x, const std::shared_ptr& norm) { + auto c = ctx->ggml_ctx; + // [B, C, T, H, W] -> [B*T, C, H, W]: GroupNorm never mixes frames. + x = ggml_cont(c, ggml_permute(c, x, 0, 1, 3, 2)); + x = std::dynamic_pointer_cast(norm)->forward(ctx, x); + return ggml_cont(c, ggml_permute(c, x, 0, 1, 3, 2)); + } + + class ResnetBlock3D : public UnaryBlock { + public: + ResnetBlock3D(int in, int out) { + blocks["norm1"] = std::make_shared(in); + blocks["norm2"] = std::make_shared(out); + blocks["conv1"] = std::make_shared(in, out); + blocks["conv2"] = std::make_shared(out, out); + if (in != out) + blocks["conv_shortcut"] = std::make_shared(in, out, std::tuple{1, 1, 1}, std::tuple{1, 1, 1}, std::tuple{0, 0, 0}); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto h = ggml_silu(ctx->ggml_ctx, causal_norm(ctx, x, blocks["norm1"])); + h = std::dynamic_pointer_cast(blocks["conv1"])->forward(ctx, h); + h = ggml_silu(ctx->ggml_ctx, causal_norm(ctx, h, blocks["norm2"])); + h = std::dynamic_pointer_cast(blocks["conv2"])->forward(ctx, h); + if (blocks.count("conv_shortcut")) + x = std::dynamic_pointer_cast(blocks["conv_shortcut"])->forward(ctx, x); + return ggml_add(ctx->ggml_ctx, h, x); + } + }; + + class SpatialAttention : public UnaryBlock { + public: + explicit SpatialAttention(int channels) { + blocks["group_norm"] = std::make_shared(channels); + for (const auto& name : {"to_q", "to_k", "to_v", "to_out.0"}) + blocks[name] = std::make_shared(channels, channels); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto c = ctx->ggml_ctx; + const int64_t width = x->ne[0], height = x->ne[1], frames = x->ne[2], channels = x->ne[3]; + auto h = causal_norm(ctx, x, blocks["group_norm"]); + h = ggml_reshape_3d(c, h, width * height, frames, channels); + h = ggml_cont(c, ggml_permute(c, h, 1, 2, 0, 3)); + auto linear = [&](const char* name, ggml_tensor* value) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, value); }; + h = ggml_ext_attention_ext(c, ctx->backend, linear("to_q", h), linear("to_k", h), linear("to_v", h), 1, nullptr, false, ctx->flash_attn_enabled); + h = linear("to_out.0", h); + h = ggml_cont(c, ggml_permute(c, h, 2, 0, 1, 3)); + return ggml_add(c, x, ggml_reshape_4d(c, h, width, height, frames, channels)); + } + }; + + class Upsample3D : public UnaryBlock { + int temporal_; + + public: + Upsample3D(int channels, bool temporal) + : temporal_(temporal ? 2 : 1) { + blocks["upscale_conv"] = std::make_shared(channels, channels * 4 * temporal_, std::tuple{1, 1, 1}, std::tuple{1, 1, 1}, std::tuple{0, 0, 0}); + blocks["conv"] = std::make_shared(channels, channels); + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + auto c = ctx->ggml_ctx; + const int64_t w = x->ne[0], h = x->ne[1], t = x->ne[2], ch = x->ne[3]; + x = std::dynamic_pointer_cast(blocks["upscale_conv"])->forward(ctx, x); + // [B, (y x z C), T, H, W] -> [B, C, T*z, H*y, W*x]. + x = ggml_reshape_4d(c, x, w * h * t, ch, temporal_, 4); + x = ggml_cont(c, ggml_permute(c, x, 1, 2, 0, 3)); + x = ggml_reshape_4d(c, x, temporal_, w * h, t, ch * 4); + x = ggml_cont(c, ggml_permute(c, x, 1, 0, 2, 3)); + x = ggml_reshape_4d(c, x, w, h * t * temporal_ * ch, 2, 2); + x = ggml_cont(c, ggml_permute(c, x, 1, 2, 0, 3)); + x = ggml_reshape_4d(c, x, w * 2, h, t * temporal_ * ch, 2); + x = ggml_cont(c, ggml_permute(c, x, 0, 2, 3, 1)); + x = ggml_reshape_4d(c, x, w * 2, h * 2, t * temporal_, ch); + if (temporal_ == 2) { + auto first = ggml_ext_slice(c, x, 2, 0, 1); + x = t == 1 ? first : ggml_concat(c, first, ggml_ext_slice(c, x, 2, 2, t * 2), 2); + } + return std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); + } + }; + + class VideoAutoencoderKL : public GGMLBlock { + SeedVR2VAEConfig config_; + + public: + explicit VideoAutoencoderKL(const SeedVR2VAEConfig& config) + : config_(config) { + for (bool decode : {false, true}) { + std::string root = decode ? "decoder." : "encoder."; + auto channels = config.channels; + if (decode) + std::reverse(channels.begin(), channels.end()); + int previous = channels.front(); + blocks[root + "conv_in"] = std::make_shared(decode ? 16 : 3, previous); + for (int i = 0; i < 4; ++i) { + auto stage = root + (decode ? "up_blocks." : "down_blocks.") + std::to_string(i) + "."; + for (int j = 0; j < config.layers + int(decode); ++j) { + blocks[stage + "resnets." + std::to_string(j)] = std::make_shared(previous, channels[i]); + previous = channels[i]; + } + if (i < 3) { + if (decode) + blocks[stage + "upsamplers.0"] = std::make_shared(previous, i < 2); + else + blocks[stage + "downsamplers.0.conv"] = std::make_shared(previous, previous, + std::tuple{i ? 3 : 1, 3, 3}, std::tuple{i ? 2 : 1, 2, 2}, std::tuple{i ? 1 : 0, 0, 0}); + } + } + blocks[root + "mid_block.resnets.0"] = std::make_shared(512, 512); + blocks[root + "mid_block.attentions.0"] = std::make_shared(512); + blocks[root + "mid_block.resnets.1"] = std::make_shared(512, 512); + blocks[root + "conv_norm_out"] = std::make_shared(previous); + blocks[root + "conv_out"] = std::make_shared(previous, decode ? 3 : 32); + } + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, bool decode) { + auto c = ctx->ggml_ctx; + std::string root = decode ? "decoder." : "encoder."; + auto apply = [&](const std::string& name) { + x = std::dynamic_pointer_cast(blocks[root + name])->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(x, "seedvr2.vae." + root + name, "hidden"); + }; + auto mid = [&]() { + apply("mid_block.resnets.0"); + apply("mid_block.attentions.0"); + apply("mid_block.resnets.1"); + }; + apply("conv_in"); + if (decode) + mid(); + for (int i = 0; i < 4; ++i) { + auto stage = std::string(decode ? "up_blocks." : "down_blocks.") + std::to_string(i) + "."; + for (int j = 0; j < config_.layers + int(decode); ++j) + apply(stage + "resnets." + std::to_string(j)); + if (i < 3) { + if (!decode) + x = ggml_pad(c, x, 1, 1, 0, 0); + apply(stage + (decode ? "upsamplers.0" : "downsamplers.0.conv")); + } + } + if (!decode) + mid(); + x = ggml_silu(c, causal_norm(ctx, x, blocks[root + "conv_norm_out"])); + apply("conv_out"); + return decode ? x : ggml_ext_slice(c, x, 3, 0, 16); + } + }; + + struct SeedVR2VAERunner : GGMLRunner { + SeedVR2VAEConfig config; + VideoAutoencoderKL model; + + SeedVR2VAERunner(ggml_backend_t backend, const String2TensorStorage& weights, std::shared_ptr manager) + : GGMLRunner(backend, manager), config(SeedVR2VAEConfig::detect_from_weights(weights, "vae")), model(config) { + model.init(params_ctx, weights, "vae"); + } + std::string get_desc() override { return "SeedVR2 VAE"; } + void get_param_tensors(std::map& tensors) { model.get_param_tensors(tensors, "vae"); } + sd::Tensor compute(int n_threads, const sd::Tensor& input, bool decode) { + auto graph = [&]() { + auto gf = new_graph_custom(MAX_GRAPH_SIZE * 4); + auto ctx = get_context(gf); + auto out = model.forward(&ctx, make_input(input), decode); + ggml_build_forward_expand(gf, out); + return gf; + }; + auto output = GGMLRunner::compute(graph, n_threads); + return restore_trailing_singleton_dims(output, 4); + } + }; +} + +#endif // __SD_MODEL_VAE_SEEDVR2_VAE_HPP__ diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 78a5e1061a..738723fa6b 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1161,31 +1161,31 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) { {"mid.attn_1.", "mid_block.attentions.0."}, }; - static std::vector> vae_conversion_map_layer; - if (vae_conversion_map_layer.empty()) { + static const auto vae_conversion_map_layer = [] { + std::vector> mappings; for (int i = 0; i < 4; ++i) { // --- encoder down blocks --- for (int j = 0; j < 2; ++j) { std::string hf_down_prefix = "encoder.down_blocks." + std::to_string(i) + ".resnets." + std::to_string(j) + "."; std::string sd_down_prefix = "encoder.down." + std::to_string(i) + ".block." + std::to_string(j) + "."; - vae_conversion_map_layer.emplace_back(sd_down_prefix, hf_down_prefix); + mappings.emplace_back(sd_down_prefix, hf_down_prefix); } if (i < 3) { std::string hf_downsample_prefix = "down_blocks." + std::to_string(i) + ".downsamplers.0."; std::string sd_downsample_prefix = "down." + std::to_string(i) + ".downsample."; - vae_conversion_map_layer.emplace_back(sd_downsample_prefix, hf_downsample_prefix); + mappings.emplace_back(sd_downsample_prefix, hf_downsample_prefix); std::string hf_upsample_prefix = "up_blocks." + std::to_string(i) + ".upsamplers.0."; std::string sd_upsample_prefix = "up." + std::to_string(3 - i) + ".upsample."; - vae_conversion_map_layer.emplace_back(sd_upsample_prefix, hf_upsample_prefix); + mappings.emplace_back(sd_upsample_prefix, hf_upsample_prefix); } // --- decoder up blocks (reverse) --- for (int j = 0; j < 3; ++j) { std::string hf_up_prefix = "decoder.up_blocks." + std::to_string(i) + ".resnets." + std::to_string(j) + "."; std::string sd_up_prefix = "decoder.up." + std::to_string(3 - i) + ".block." + std::to_string(j) + "."; - vae_conversion_map_layer.emplace_back(sd_up_prefix, hf_up_prefix); + mappings.emplace_back(sd_up_prefix, hf_up_prefix); } } @@ -1193,9 +1193,10 @@ std::string convert_diffusers_vae_to_original_sd1(std::string name) { for (int i = 0; i < 2; ++i) { std::string hf_mid_res_prefix = "mid_block.resnets." + std::to_string(i) + "."; std::string sd_mid_res_prefix = "mid.block_" + std::to_string(i + 1) + "."; - vae_conversion_map_layer.emplace_back(sd_mid_res_prefix, hf_mid_res_prefix); + mappings.emplace_back(sd_mid_res_prefix, hf_mid_res_prefix); } - } + return mappings; + }(); static const std::vector> vae_conversion_map_attn = { {"norm.", "group_norm."}, diff --git a/src/pipeline/relay_vsr.cpp b/src/pipeline/relay_vsr.cpp index bc5cad7189..c8b0d72fdf 100644 --- a/src/pipeline/relay_vsr.cpp +++ b/src/pipeline/relay_vsr.cpp @@ -84,7 +84,7 @@ namespace sd { sources.emplace(source->second, tensor); ggml_set_name(tensor, source->second.c_str()); } - runner->set_flash_attention_enabled(true); + runtime.configure_runner(*runner, module, true, module == SDBackendModule::DIFFUSION); return runtime.manager->register_param_tensors(component, std::move(sources), runtime.backends.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, compute, params); diff --git a/src/pipeline/seedvr2.cpp b/src/pipeline/seedvr2.cpp new file mode 100644 index 0000000000..5fe2eaa444 --- /dev/null +++ b/src/pipeline/seedvr2.cpp @@ -0,0 +1,167 @@ +#include "pipeline/video_upscale.h" + +#include + +#include "core/rng_philox.hpp" +#include "generation_info.h" +#include "model/diffusion/seedvr2.hpp" +#include "model/vae/seedvr2_vae.hpp" + +namespace sd { + namespace pipeline { + + class SeedVR2Upscaler : public VideoUpscaleModel { + std::unique_ptr dit_; + std::unique_ptr vae_; + int scale_ = 0; + int n_threads_ = 1; + + public: + bool init(VideoUpscaleRuntime& runtime, const String2TensorStorage& weights, int upscale_factor) override { + scale_ = upscale_factor; + n_threads_ = runtime.n_threads; + dit_ = std::make_unique(runtime.backends.runtime_backend(SDBackendModule::DIFFUSION), weights, runtime.manager); + vae_ = std::make_unique(runtime.backends.runtime_backend(SDBackendModule::VAE), weights, runtime.manager); + std::set used; + auto register_model = [&](auto& runner, ModelComponent component, SDBackendModule module) { + std::map tensors; + runner->get_param_tensors(tensors); + for (const auto& [name, tensor] : tensors) { + if (!weights.count(name)) + throw std::runtime_error("Missing SeedVR2 tensor: " + name); + used.insert(name); + } + runtime.configure_runner(*runner, module, module == SDBackendModule::VAE, module == SDBackendModule::VAE); + return runtime.manager->register_param_tensors(component, std::move(tensors), + runtime.backends.params_backend_is_disk(module) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, + runtime.backends.runtime_backend(module), runtime.backends.params_backend(module)); + }; + if (!register_model(dit_, ModelComponent::Diffusion, SDBackendModule::DIFFUSION) || + !register_model(vae_, ModelComponent::VAE, SDBackendModule::VAE) || !runtime.manager->validate_registered_tensors()) + return false; + for (const auto& [name, storage] : weights) { + if (!used.count(name) && name != "diffusion.negative_conditioning" && !ends_with(name, ".comfy_quant")) + throw std::runtime_error("Unused SeedVR2 tensor: " + name); + } + LOG_INFO("SeedVR2 validated %zu parameter tensors; negative conditioning is unused at guidance 1", used.size()); + return true; + } + + void end() override { + if (dit_) + dit_->runner_end(); + if (vae_) + vae_->runner_end(); + } + + bool upscale(const sd_video_upscale_params_t& params, sd_video_generation_result_t* result) override { + const int chunk = params.temporal_chunk_size; + if (chunk < 1 || chunk > INT_MAX - 3 || (chunk - 1) % 4) { + LOG_ERROR("SeedVR2 temporal chunk size must be 1 or 4n+1 input frames"); + return false; + } + struct Reset { + SeedVR2Upscaler* model; + ~Reset() { model->end(); } + } reset{this}; + end(); + const int width = int(params.frames[0].width) * scale_, height = int(params.frames[0].height) * scale_; + if (width > INT_MAX - 15 || height > INT_MAX - 15) + return false; + const int pw = (width + 15) / 16 * 16, ph = (height + 15) / 16 * 16; + PhiloxRNG rng; + rng.manual_seed(uint64_t(params.seed)); + for (int start = 0; start < params.frame_count;) { + const int count = std::min(chunk, params.frame_count - start); + const int padded_count = count == 1 ? 1 : ((count - 1 + 3) / 4) * 4 + 1; + LOG_INFO("SeedVR2 chunk %d..%d/%d", start + 1, start + count, params.frame_count); + const int64_t begin = ggml_time_us(); + std::vector> images; + for (int i = 0; i < count; ++i) { + auto image = sd_image_to_tensor(params.frames[start + i], -1, -1, false); + image /= 255.f; + image = sd::ops::interpolate(image, {width, height, 3, 1}, sd::ops::InterpolateMode::Bicubic); + image = sd::ops::clamp(image, 0.f, 1.f); + if (pw != width) + image = sd::ops::concat(image, sd::zeros({pw - width, height, 3, 1}), 0); + if (ph != height) + image = sd::ops::concat(image, sd::zeros({pw, ph - height, 3, 1}), 1); + image *= 2.f; + image -= 1.f; + images.push_back(image); + } + auto pixels = images.front(); + for (int i = 1; i < padded_count; ++i) + pixels = sd::ops::concat(pixels, images[std::min(i, count - 1)], 3); + pixels.permute_({0, 1, 3, 2}); + const int64_t prepared = ggml_time_us(); + auto latent = vae_->compute(n_threads_, pixels, false); + if (latent.empty()) + return false; + latent *= 0.9152f; + const int64_t encoded = ggml_time_us(); + const int lw = pw / 8, lh = ph / 8, lt = (padded_count - 1) / 4 + 1; + sd::Tensor noise({lw, lh, lt, 16}); + auto values = rng.randn(uint32_t(noise.numel())); + std::copy(values.begin(), values.end(), noise.data()); + auto mask = sd::zeros({lw, lh, lt, 1}); + mask += 1.f; + auto input = sd::ops::concat(sd::ops::concat(noise, latent, 3), mask, 3); + input.reshape_({2, lw / 2, 2, lh / 2, lt, 33}); + input.permute_({5, 0, 2, 1, 3, 4}); + input.reshape_({132, (lw / 2) * (lh / 2) * lt}); + auto prediction = dit_->compute(n_threads_, input, lw / 2, lh / 2, lt); + if (prediction.empty()) + return false; + prediction.reshape_({16, 2, 2, lw / 2, lh / 2, lt}); + prediction.permute_({1, 3, 2, 4, 5, 0}); + prediction.reshape_({lw, lh, lt, 16}); + latent = (noise - prediction) / 0.9152f; + const int64_t denoised = ggml_time_us(); + auto decoded = vae_->compute(n_threads_, latent, true); + if (decoded.empty()) + return false; + const int64_t decoded_at = ggml_time_us(); + if (!std::all_of(decoded.values().begin(), decoded.values().end(), [](float x) { return std::isfinite(x); })) + throw std::runtime_error("SeedVR2 produced non-finite pixels"); + decoded = sd::ops::clamp((decoded + 1.f) * 127.5f, 0.f, 255.f); + for (int i = 0; i < count; ++i) { + auto output = sd::ops::slice(sd::ops::slice(sd::ops::slice(decoded, 2, i, i + 1), 0, 0, width), 1, 0, height); + output.permute_({3, 0, 1, 2}); + auto& image = result->frames[start + i]; + image.width = width; + image.height = height; + image.channel = 3; + image.data = static_cast(malloc(size_t(width) * height * 3)); + if (!image.data) + throw std::bad_alloc(); + std::transform(output.values().begin(), output.values().end(), image.data, [](float x) { return uint8_t(x); }); + } + if (auto progress = sd_get_progress_callback()) + progress(start + count, params.frame_count, 0.f, sd_get_progress_callback_data()); + LOG_INFO("SeedVR2 seconds: preprocess %.3f, VAE encode %.3f, DiT and latent transforms %.3f, VAE decode %.3f, postprocess %.3f", + (prepared - begin) / 1e6, (encoded - prepared) / 1e6, (denoised - encoded) / 1e6, (decoded_at - denoised) / 1e6, (ggml_time_us() - decoded_at) / 1e6); + start += count; + } + auto info = sd::upscale_generation_info(scale_); + info["kind"] = "video"; + info["upscale"]["model"] = "seedvr2"; + info["upscale"]["temporal_chunk_size"] = chunk; + info["seed"] = params.seed; + info["fps"] = params.fps; + info["video_frames"] = params.frame_count; + result->generation_info = sd::copy_generation_info(info); + return result->generation_info != nullptr; + } + }; + + bool is_seedvr2_model(const String2TensorStorage& weights) { + return weights.count("diffusion.vid_in.proj.weight") && weights.count("diffusion.blocks.0.attn.rope.rope.freqs") && + weights.count("diffusion.positive_conditioning"); + } + + std::unique_ptr create_seedvr2_upscaler() { + return std::make_unique(); + } + } +} diff --git a/src/pipeline/video_upscale.cpp b/src/pipeline/video_upscale.cpp index 5b34a9288c..f9b631e637 100644 --- a/src/pipeline/video_upscale.cpp +++ b/src/pipeline/video_upscale.cpp @@ -2,11 +2,62 @@ #include +#include "core/ggml_runner.h" +#include "core/probe/conv_direct.h" +#include "core/probe/flash_attn.h" #include "core/util.h" namespace sd { namespace pipeline { + bool VideoUpscaleRuntime::flash_attention_for(SDBackendModule module) { + const auto mode = flash_attn_assignment.get(module); + if (mode == "on" || mode == "off") + return mode == "on"; + const auto devices = backends.runtime_backends(module); + bool enabled = !devices.empty(), any_enabled = false; + for (auto backend : devices) { + auto it = flash_attn_by_backend.find(backend); + if (it == flash_attn_by_backend.end()) + it = flash_attn_by_backend.emplace(backend, sd_probe_flash_attention(backend, n_threads)).first; + enabled = enabled && it->second; + any_enabled = any_enabled || it->second; + } + if (any_enabled && !enabled) + LOG_INFO("flash attention auto: %s -> off (compute devices disagree)", sd_backend_module_name(module)); + return enabled; + } + + bool VideoUpscaleRuntime::conv_direct_for(SDBackendModule module, bool conv3d) { + const auto mode = conv_direct_assignment.get(module); + if (mode == "on" || mode == "off") + return mode == "on"; + auto& cache = conv3d ? conv3d_direct_by_backend : conv2d_direct_by_backend; + const auto devices = backends.runtime_backends(module); + bool enabled = !devices.empty(), any_enabled = false; + for (auto backend : devices) { + auto it = cache.find(backend); + if (it == cache.end()) + it = cache.emplace(backend, sd_probe_conv_direct(backend, conv3d)).first; + enabled = enabled && it->second; + any_enabled = any_enabled || it->second; + } + if (any_enabled && !enabled) + LOG_INFO("conv%dd direct auto: %s -> off (compute devices disagree)", conv3d ? 3 : 2, sd_backend_module_name(module)); + return enabled; + } + + void VideoUpscaleRuntime::configure_runner(GGMLRunner& runner, SDBackendModule module, bool conv2d, bool conv3d) { + const bool flash = flash_attention_for(module); + const bool direct2d = conv2d && conv_direct_for(module, false); + const bool direct3d = conv3d && conv_direct_for(module, true); + runner.set_flash_attention_enabled(flash); + runner.set_conv2d_direct_enabled(direct2d); + runner.set_conv3d_direct_enabled(direct3d); + LOG_INFO("video upscale %s: flash attention %s, conv2d direct %s, conv3d direct %s", + sd_backend_module_name(module), flash ? "on" : "off", direct2d ? "on" : "off", direct3d ? "on" : "off"); + } + VideoUpscaler::~VideoUpscaler() { if (model_) model_->end(); @@ -26,6 +77,11 @@ namespace sd { executor_ = std::make_unique(runtime_.n_threads); sd::ParallelScope scope(executor_.get()); std::string error; + if (!sd_parse_flash_attention_assignment(params.flash_attn_mode ? params.flash_attn_mode : "", &runtime_.flash_attn_assignment, &error) || + !sd_parse_conv_direct_assignment(params.conv_direct_mode ? params.conv_direct_mode : "", &runtime_.conv_direct_assignment, &error)) { + LOG_ERROR("video upscale operator configuration: %s", error.c_str()); + return false; + } if (!runtime_.backends.init(params.backend, params.params_backend, nullptr, &error)) { LOG_ERROR("video upscaler backend: %s", error.c_str()); return false; @@ -41,6 +97,8 @@ namespace sd { const auto& weights = loader.get_tensor_storage_map(); if (is_relay_vsr_model(weights)) model_ = create_relay_vsr_upscaler(); + else if (is_seedvr2_model(weights)) + model_ = create_seedvr2_upscaler(); else { LOG_ERROR("unsupported video upscale model architecture"); return false; @@ -104,10 +162,11 @@ void sd_video_upscaler_ctx_params_init(sd_video_upscaler_ctx_params_t* params) { void sd_video_upscale_params_init(sd_video_upscale_params_t* params) { if (!params) return; - *params = {}; - params->fps = 24; - params->seed = 42; - params->window_size = 16; + *params = {}; + params->fps = 24; + params->seed = 42; + params->window_size = 16; + params->temporal_chunk_size = 5; } video_upscaler_ctx_t* new_video_upscaler_ctx(const sd_video_upscaler_ctx_params_t* params) { diff --git a/src/pipeline/video_upscale.h b/src/pipeline/video_upscale.h index 62d0af1039..ec5815c44f 100644 --- a/src/pipeline/video_upscale.h +++ b/src/pipeline/video_upscale.h @@ -1,6 +1,7 @@ #ifndef __SD_PIPELINE_VIDEO_UPSCALE_H__ #define __SD_PIPELINE_VIDEO_UPSCALE_H__ +#include #include #include "core/ggml_extend_backend.h" @@ -8,6 +9,8 @@ #include "model_manager.h" #include "stable-diffusion.h" +struct GGMLRunner; + namespace sd { namespace pipeline { @@ -15,6 +18,15 @@ namespace sd { SDBackendManager backends; std::shared_ptr manager; int n_threads = 1; + SDBackendAssignment flash_attn_assignment; + SDBackendAssignment conv_direct_assignment; + std::map flash_attn_by_backend; + std::map conv2d_direct_by_backend; + std::map conv3d_direct_by_backend; + + bool flash_attention_for(SDBackendModule module); + bool conv_direct_for(SDBackendModule module, bool conv3d); + void configure_runner(GGMLRunner& runner, SDBackendModule module, bool conv2d, bool conv3d); }; class VideoUpscaleModel { @@ -27,6 +39,8 @@ namespace sd { bool is_relay_vsr_model(const String2TensorStorage& weights); std::unique_ptr create_relay_vsr_upscaler(); + bool is_seedvr2_model(const String2TensorStorage& weights); + std::unique_ptr create_seedvr2_upscaler(); class VideoUpscaler { VideoUpscaleRuntime runtime_; From b43ce4badd8130e743be3096024f4fac5b05719e Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 02:11:28 +0800 Subject: [PATCH 36/57] fix: preserve Krea2 text-fusion names in Kohya LoRAs (#2301) --- src/name_conversion.cpp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 738723fa6b..6bcd483e31 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1403,8 +1403,8 @@ std::string convert_pmid_v2_name(const std::string& name) { return name; } -std::string convert_sep_to_dot(std::string name) { - const std::vector protected_tokens = { +std::string convert_sep_to_dot(std::string name, SDVersion version) { + std::vector protected_tokens = { "self_attn", "out_proj", "q_proj", @@ -1480,6 +1480,11 @@ std::string convert_sep_to_dot(std::string name) { "output_proj", "token_refiner", }; + if (sd_version_is_krea2(version)) { + protected_tokens.push_back("text_fusion"); + protected_tokens.push_back("layerwise_blocks"); + protected_tokens.push_back("refiner_blocks"); + } // record the positions of underscores that should NOT be replaced std::unordered_set protected_positions; @@ -1729,7 +1734,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { // LOG_VERBOSE("name %s %d", name.c_str(), version); if (sd_version_is_unet(version) || is_underline || is_lycoris_underline) { - name = convert_sep_to_dot(name); + name = convert_sep_to_dot(name, version); } } From fc350be2ceeeb859bcc201ac435591ed7d46821d Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 02:32:27 +0800 Subject: [PATCH 37/57] feat(server): report generation stages independently of previews (#2303) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Maxime DrΓ©an --- examples/server/README.md | 10 ++++ examples/server/api.md | 52 +++++++++++++++++ examples/server/async_jobs.cpp | 96 ++++++++++++++++++++++++++++--- examples/server/async_jobs.h | 14 +++++ examples/server/frontend | 2 +- include/stable-diffusion.h | 14 +++++ src/pipeline/diffusion_engine.cpp | 8 +++ src/pipeline/diffusion_engine.h | 5 ++ src/pipeline/image.cpp | 9 ++- src/pipeline/video.cpp | 8 ++- src/runtime/generation_progress.h | 42 ++++++++++++++ src/stable-diffusion.cpp | 10 ++++ 12 files changed, 259 insertions(+), 11 deletions(-) create mode 100644 src/runtime/generation_progress.h diff --git a/examples/server/README.md b/examples/server/README.md index 8d8f687cb5..d93d829bc7 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -76,6 +76,16 @@ Recipes declaring `video_frames` also switch the form to Video Generation; other recipes switch to Image Generation. Generation is enabled after the selected recipe is loaded. Changing a selection does not reload the model until you click Load. + +The generation status bar shows the current stage and elapsed time even when +Preview Mode is `none`. During sampling it shows completed steps for the current +pass; prompt encoding and VAE decoding show an activity indicator when their +total work is unknown. Sampling reaching 100% does not mean the job is finished: +decoding, ADetailer, and output encoding may still follow. +The canvas keeps the submitted output aspect ratio, and stage changes keep the +same layout. Editing the next request's resolution does not resize the current +preview or result. + The server can also start with the existing model arguments shown below; existing API clients then work without selecting a recipe. The UI automatically selects **Command-line model (all features)** and uses the startup generation diff --git a/examples/server/api.md b/examples/server/api.md index 9c8d8c8ce9..f6096c67a1 100644 --- a/examples/server/api.md +++ b/examples/server/api.md @@ -775,6 +775,7 @@ Field types: | `queue_position` | `integer` | | `result` | `object \| null` | | `error` | `object \| null` | +| `progress` | `object` | | `preview` | `object \| null` | `created`, `started`, and `completed` are integer Unix timestamps in @@ -782,6 +783,57 @@ Field types: the elapsed generation time in seconds, including output encoding and excluding queue time. +`progress` is available independently of the preview setting, for both image +and video jobs. Continue polling until `status` is terminal: a sampling pass +can finish before decoding, ADetailer, or output encoding finishes. + +| Field | Type | Notes | +| --- | --- | --- | +| `progress.stage` | `string` | Current execution stage; see below | +| `progress.detail` | `string` | Stage detail, or an empty string | +| `progress.pass` | `integer` | Sampling pass number within the job; zero before sampling | +| `progress.completed` | `integer \| null` | Completed units in the current stage/pass | +| `progress.total` | `integer \| null` | Actual total units, or null when unknown | +| `progress.unit` | `string \| null` | `steps` for sampling, `images` for image output encoding; null when uncounted | +| `progress.stage_elapsed_ms` | `integer` | Elapsed time in the current stage/detail/pass | +| `progress.elapsed_ms` | `integer` | Job execution time, excluding queue time; zero before execution | + +Stages include `queued`, `waiting` (waiting for the execution lock), `preparing`, +`conditioning`, `sampling`, `vae_encode`, `vae_decode`, `audio_decode`, +`upscaling`, `adetailer`, and `encoding`. Stages may repeat or be skipped. +Details include `positive_prompt` / `negative_prompt` / `image_condition` for +conditioning, `high_noise` / `low_noise` for two-stage sampling, and `image` / +`video` for output encoding. During ADetailer, `detail` identifies the inner +stage, such as `conditioning` or `sampling`. Clients should tolerate unknown +stage, detail, and unit values. + +Sampling counts completed logical steps using the actual schedule, including +strength adjustments. Each new sampling pass increments `pass` and starts at +zero, including batch items, highres passes, high/low-noise passes, and ADetailer +passes. The ratio is **current-pass progress**, not overall job completion. +The progress snapshot is published atomically. Preview metadata still describes +the preview frame, which may be older than the current progress. + +Unknown totals are represented by null counts, not an estimated percentage. +For example, prompt encoding can report: + +```json +"progress": { + "stage": "conditioning", + "detail": "positive_prompt", + "pass": 0, + "completed": null, + "total": null, + "unit": null, + "stage_elapsed_ms": 3200, + "elapsed_ms": 4100 +} +``` + +Elapsed times freeze when a job completes, fails, or is cancelled; the last +stage is retained for inspection. Use `status` to determine the final outcome. +Elapsed time alone does not indicate that computation is advancing. + `preview` sub-fields: | Field | Type | Notes | diff --git a/examples/server/async_jobs.cpp b/examples/server/async_jobs.cpp index 7a80107897..c730b9c13d 100644 --- a/examples/server/async_jobs.cpp +++ b/examples/server/async_jobs.cpp @@ -92,6 +92,63 @@ static void set_preview_callback(AsyncGenerationJob& job) { sd_set_preview_callback(preview_callback, mode, interval, true, false, &job); } +static void update_job_progress(AsyncGenerationJob& job, const char* stage, const char* detail = "", int sample_pass = -1, int64_t completed = 0, int64_t total = 0, const char* unit = "") { + std::lock_guard lock(job.progress_mutex); + auto& progress = job.progress; + if (sample_pass < 0) { + sample_pass = progress.sample_pass; + } + if (progress.stage != stage || progress.detail != detail || progress.sample_pass != sample_pass) { + progress.stage_started = std::chrono::steady_clock::now(); + } + progress.stage = stage; + progress.detail = detail; + progress.sample_pass = sample_pass; + progress.completed = completed; + progress.total = total; + progress.unit = unit; +} + +static void finish_job_progress(AsyncGenerationJob& job) { + std::lock_guard lock(job.progress_mutex); + job.progress.finished = std::chrono::steady_clock::now(); +} + +class JobCallbacksGuard { + sd_ctx_t* ctx_; + AsyncGenerationJob& job_; + bool detailing_ = false; + + static void progress_callback(const sd_generation_progress_t* progress, void* data) { + auto& guard = *static_cast(data); + update_job_progress(guard.job_, guard.detailing_ ? "adetailer" : progress->stage, + guard.detailing_ ? progress->stage : progress->detail, + progress->sample_pass, progress->completed, progress->total, progress->unit); + } + +public: + JobCallbacksGuard(sd_ctx_t* ctx, AsyncGenerationJob& job) + : ctx_(ctx), job_(job) { + clear_preview_callback(); + set_preview_callback(job); + sd_set_generation_progress_callback(ctx_, progress_callback, this); + } + + ~JobCallbacksGuard() { + sd_set_generation_progress_callback(ctx_, nullptr, nullptr); + clear_preview_callback(); + } + + void begin_detailing() { + clear_preview_callback(); + detailing_ = true; + update_job_progress(job_, "adetailer"); + } + + JobCallbacksGuard(const JobCallbacksGuard&) = delete; + JobCallbacksGuard& operator=(const JobCallbacksGuard&) = delete; +}; + const char* async_job_kind_name(AsyncJobKind kind) { switch (kind) { case AsyncJobKind::ImgGen: @@ -186,6 +243,7 @@ bool cancel_queued_job(AsyncJobManager& manager, AsyncGenerationJob& job) { job.result_fps = 0; job.error_code = "cancelled"; job.error_message = "job cancelled by client"; + finish_job_progress(job); return true; } @@ -200,6 +258,24 @@ json make_async_job_json(const AsyncJobManager& manager, const AsyncGenerationJo result["completed"] = job.completed_at == 0 ? json(nullptr) : json(job.completed_at); result["queue_position"] = 0; + { + std::lock_guard lock(job.progress_mutex); + const auto& progress = job.progress; + const auto end = progress.finished == std::chrono::steady_clock::time_point{} + ? std::chrono::steady_clock::now() + : progress.finished; + result["progress"] = { + {"stage", progress.stage}, + {"detail", progress.detail}, + {"pass", progress.sample_pass}, + {"completed", progress.total > 0 ? json(progress.completed) : json(nullptr)}, + {"total", progress.total > 0 ? json(progress.total) : json(nullptr)}, + {"unit", progress.unit.empty() ? json(nullptr) : json(progress.unit)}, + {"stage_elapsed_ms", std::chrono::duration_cast(end - progress.stage_started).count()}, + {"elapsed_ms", job.started_at == 0 ? 0 : std::max(0, (job.completed_at == 0 ? unix_timestamp_now_ms() : job.completed_at) - job.started_at)}, + }; + } + if (job.status == AsyncJobStatus::Queued) { size_t position = 1; for (const auto& queued_id : manager.queue) { @@ -278,17 +354,19 @@ bool execute_img_gen_job(ServerRuntime& runtime, return false; } JobCancellationMonitor cancellation(runtime.sd_ctx, job); - set_preview_callback(job); + JobCallbacksGuard callbacks(runtime.sd_ctx, job); sd_image_generation_result_t* generated = nullptr; generate_image(runtime.sd_ctx, ¶ms, &generated); results.adopt(generated, runtime.ctx_params->lora_model_dir); - clear_preview_callback(); if (job.cancel_requested.load()) { return false; } - if (!results.empty() && !apply_server_adetailer(runtime, job.img_gen, params, results)) { - error_message = "ADetailer failed"; - return false; + if (!results.empty() && job.img_gen.adetailer_enabled && !runtime.default_gen_params->ad_model_path.empty()) { + callbacks.begin_detailing(); + if (!apply_server_adetailer(runtime, job.img_gen, params, results)) { + error_message = "ADetailer failed"; + return false; + } } } @@ -305,6 +383,7 @@ bool execute_img_gen_job(ServerRuntime& runtime, encoded_format = EncodedImageFormat::WEBP; } + update_job_progress(job, "encoding", "image", -1, 0, num_results, "images"); for (int i = 0; i < num_results; ++i) { if (results[i].data == nullptr) { continue; @@ -325,6 +404,7 @@ bool execute_img_gen_job(ServerRuntime& runtime, } output_images.push_back(base64_encode(image_bytes)); output_info.push_back(json::parse(results.generation_info(i))); + update_job_progress(job, "encoding", "image", -1, i + 1, num_results, "images"); } if (output_images.empty()) { @@ -356,7 +436,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime, return false; } JobCancellationMonitor cancellation(runtime.sd_ctx, job); - set_preview_callback(job); + JobCallbacksGuard callbacks(runtime.sd_ctx, job); sd_video_generation_result_t* generated = nullptr; if (generate_video(runtime.sd_ctx, ¶ms, &generated)) { output_fps = generated->fps; @@ -367,7 +447,6 @@ bool execute_vid_gen_job(ServerRuntime& runtime, if (job.vid_gen.gen_params.embed_image_metadata) str_params = results.generation_info(0); } - clear_preview_callback(); } num_results = results.count(); @@ -377,6 +456,7 @@ bool execute_vid_gen_job(ServerRuntime& runtime, return false; } + update_job_progress(job, "encoding", "video"); std::vector video_bytes = create_video_from_sd_images_to_vector(job.vid_gen.output_format, results.data(), num_results, @@ -425,6 +505,7 @@ void async_job_worker(ServerRuntime& runtime) { job = it->second; job->status = AsyncJobStatus::Generating; job->started_at = unix_timestamp_now_ms(); + update_job_progress(*job, "waiting"); } std::vector output_images; @@ -459,6 +540,7 @@ void async_job_worker(ServerRuntime& runtime) { } job->completed_at = unix_timestamp_now_ms(); + finish_job_progress(*job); if (job->cancel_requested.load()) { job->status = AsyncJobStatus::Cancelled; job->error_code = "cancelled"; diff --git a/examples/server/async_jobs.h b/examples/server/async_jobs.h index 6fa3cbd138..a442305606 100644 --- a/examples/server/async_jobs.h +++ b/examples/server/async_jobs.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -29,6 +30,17 @@ enum class AsyncJobStatus { const char* async_job_kind_name(AsyncJobKind kind); const char* async_job_status_name(AsyncJobStatus status); +struct AsyncJobProgress { + std::string stage = "queued"; + std::string detail; + std::string unit; + int sample_pass = 0; + int64_t completed = 0; + int64_t total = 0; + std::chrono::steady_clock::time_point stage_started = std::chrono::steady_clock::now(); + std::chrono::steady_clock::time_point finished{}; +}; + struct AsyncGenerationJob { std::string id; AsyncJobKind kind = AsyncJobKind::ImgGen; @@ -47,6 +59,8 @@ struct AsyncGenerationJob { int result_fps = 0; std::string error_code; std::string error_message; + mutable std::mutex progress_mutex; + AsyncJobProgress progress; mutable std::mutex preview_mutex; std::string preview_b64; int preview_pass = 0; diff --git a/examples/server/frontend b/examples/server/frontend index fa93429c47..d3d83992ff 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit fa93429c47c924054e6a9c14f8d7b8e06f57a0ec +Subproject commit d3d83992ff9bd54c99b11fc36771c13f8637376e diff --git a/include/stable-diffusion.h b/include/stable-diffusion.h index 53c6b7059c..7453a3f81a 100644 --- a/include/stable-diffusion.h +++ b/include/stable-diffusion.h @@ -500,6 +500,15 @@ struct ggml_tensor; typedef void (*sd_log_cb_t)(enum sd_log_level_t level, const char* text, void* data); typedef void (*sd_progress_cb_t)(int step, int steps, float time, void* data); +typedef struct { + const char* stage; + const char* detail; + const char* unit; + int sample_pass; + int64_t completed; + int64_t total; +} sd_generation_progress_t; +typedef void (*sd_generation_progress_cb_t)(const sd_generation_progress_t* progress, void* data); typedef void (*sd_preview_cb_t)(int step, int frame_count, sd_image_t* frames, bool is_noisy, void* data); typedef struct { @@ -510,6 +519,11 @@ typedef bool (*sd_graph_eval_callback_t)(struct ggml_tensor* t, bool ask, void* SD_API void sd_set_log_callback(sd_log_cb_t sd_log_cb, void* data); SD_API void sd_set_progress_callback(sd_progress_cb_t cb, void* data); +// Context-local generation events, independent of preview and the legacy progress callback. +// Strings and the event are valid only during the synchronous callback. Callbacks must not throw. +// total == 0 means unknown; sample_pass counts from 1 since callback registration. +// Register/clear only while the context is idle. Passing nullptr clears the callback. +SD_API void sd_set_generation_progress_callback(sd_ctx_t* ctx, sd_generation_progress_cb_t cb, void* data); // In each sampling pass, a positive interval previews every Nth denoiser step, while a // negative interval previews only completed logical step -interval. Zero previews the final // completed step of the first sampling pass (base-resolution or high-noise). diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 2a6683f010..1ec292c161 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -1807,6 +1807,7 @@ void StableDiffusionGGML::lora_stat() { } bool StableDiffusionGGML::apply_loras(const sd_lora_t* loras, uint32_t lora_count) { + generation_progress.report("preparing", "loras"); std::vector all_loras; all_loras.reserve(lora_count); for (uint32_t i = 0; i < lora_count; i++) { @@ -2243,6 +2244,7 @@ void StableDiffusionGGML::report_sample_progress(int step, ? (now - *last_progress_us) / 1000000.f : 0.f; pretty_progress(showstep, (int)total_steps, step_seconds); + generation_progress.report_sample(showstep, static_cast(total_steps)); if (last_progress_us != nullptr) { *last_progress_us = now; } @@ -2424,6 +2426,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr(steps)); + generation_progress.begin_sample(static_cast(steps), + work_diffusion_model == high_noise_diffusion_model ? "high_noise" : high_noise_diffusion_model != nullptr ? "low_noise" + : ""); SamplePreviewContext preview = prepare_sample_preview_context(); sd::Tensor processed_init_latent = denoiser->process_latent_in(init_latent); @@ -2995,6 +3000,7 @@ int StableDiffusionGGML::align_video_frames(int frames) { } sd::Tensor StableDiffusionGGML::encode_to_vae_latents(const sd::Tensor& x) { + generation_progress.report("vae_encode"); auto latents = first_stage_model->encode(n_threads, x, vae_tiling_params, circular_x, circular_y); if (latents.empty()) { return {}; @@ -3015,6 +3021,7 @@ sd::Tensor StableDiffusionGGML::encode_first_stage(const sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensor& x, bool decode_video) { + generation_progress.report("vae_decode"); if (sd_version_is_pid(version) || sd_version_is_minit2i(version) || version == VERSION_IRIS) { return sd::ops::clamp((x + 1.f) * 0.5f, 0.0f, 1.0f); } @@ -3069,6 +3076,7 @@ sd::Tensor StableDiffusionGGML::decode_ltx_audio_latent(const sd::Tensor< if (audio_vae_model == nullptr || audio_latent.empty()) { return {}; } + generation_progress.report("audio_decode"); auto waveform = audio_vae_model->decode(n_threads, audio_latent); return waveform; } diff --git a/src/pipeline/diffusion_engine.h b/src/pipeline/diffusion_engine.h index 97e3acac9b..990a8b6515 100644 --- a/src/pipeline/diffusion_engine.h +++ b/src/pipeline/diffusion_engine.h @@ -21,6 +21,7 @@ #include "model/adapter/lora.hpp" #include "model_builders.h" #include "model_manager.h" +#include "runtime/generation_progress.h" #include "stable-diffusion.h" class RNG; @@ -199,6 +200,7 @@ class StableDiffusionGGML { std::unique_ptr conditioning_cache_; std::vector conditioning_loras_; bool executing_ = false; + sd::GenerationProgress generation_progress; std::shared_ptr denoiser; std::vector file_alphas_cumprod; @@ -254,6 +256,9 @@ class StableDiffusionGGML { explicit ExecutionScope(StableDiffusionGGML& sd) : operation(sd) { + if (operation.acquired) { + sd.generation_progress.report("preparing"); + } ready = operation.acquired && sd.refresh_model_sources(); } diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 03a6707d40..1f7d5e4539 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -531,6 +531,7 @@ namespace sd::pipeline { condition_params.ref_image_params = ref_image_params; + sd->generation_progress.report("conditioning", "extensions"); sd->prepare_generation_extensions(request->pm_params, request->pulid_params, condition_params, @@ -538,7 +539,8 @@ namespace sd::pipeline { sd->compute_ip_adapter_tokens(sd_img_gen_params->ip_adapter_image, sd_img_gen_params->ip_adapter_strength); int64_t prepare_start_ms = ggml_time_ms(); condition_params.zero_out_masked = false; - auto cond = sd->get_learned_condition(condition_params); + sd->generation_progress.report("conditioning", "positive_prompt"); + auto cond = sd->get_learned_condition(condition_params); if (cond.empty()) { LOG_ERROR("failed to encode prompt"); return std::nullopt; @@ -584,6 +586,7 @@ namespace sd::pipeline { // Keep source latents in CFG, but omit source images from the text encoder. condition_params.ref_images = nullptr; } + sd->generation_progress.report("conditioning", "negative_prompt"); uncond = sd->get_learned_condition(condition_params); if (uncond.empty()) { LOG_ERROR("failed to encode negative prompt"); @@ -615,6 +618,7 @@ namespace sd::pipeline { if (use_ref_latent_img_cfg) { condition_params.ref_images = &empty_ref_images; } + sd->generation_progress.report("conditioning", "image_condition"); img_uncond = sd->get_learned_condition(condition_params); if (img_uncond.empty()) { LOG_ERROR("failed to encode image guidance prompt"); @@ -830,6 +834,7 @@ namespace sd::pipeline { return {}; } + sd->generation_progress.report("upscaling", "hires"); sd::Tensor upscaled_tensor; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { upscaled_tensor = upscaler->upscale_tensor(decoded); @@ -930,6 +935,7 @@ namespace sd::pipeline { ImageVaeAxesGuard axes_guard(sd, sd_img_gen_params, request); SamplePlan plan(sd, sd_img_gen_params, request); + sd->generation_progress.report("preparing", "inputs"); auto latents_opt = prepare_image_generation_latents(sd, sd_img_gen_params, &request, @@ -1027,6 +1033,7 @@ namespace sd::pipeline { } LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height); + sd->generation_progress.report("upscaling", "hires"); std::unique_ptr hires_upscaler; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { diff --git a/src/pipeline/video.cpp b/src/pipeline/video.cpp index 8c911b143a..915dc54b39 100644 --- a/src/pipeline/video.cpp +++ b/src/pipeline/video.cpp @@ -1210,7 +1210,8 @@ namespace sd::pipeline { condition_params.ref_image_params = ref_image_params; int64_t prepare_start_ms = ggml_time_ms(); - embeds.cond = sd->get_learned_condition(condition_params); + sd->generation_progress.report("conditioning", "positive_prompt"); + embeds.cond = sd->get_learned_condition(condition_params); if (embeds.cond.empty()) { LOG_ERROR("failed to encode video prompt"); return std::nullopt; @@ -1235,7 +1236,8 @@ namespace sd::pipeline { } if (request.use_uncond) { condition_params.text = request.negative_prompt; - embeds.uncond = sd->get_learned_condition(condition_params); + sd->generation_progress.report("conditioning", "negative_prompt"); + embeds.uncond = sd->get_learned_condition(condition_params); if (embeds.uncond.empty()) { LOG_ERROR("failed to encode negative video prompt"); return std::nullopt; @@ -1901,6 +1903,7 @@ namespace sd::pipeline { sd->reset_generation_extensions(); SamplePlan plan(sd, sd_vid_gen_params, request); + sd->generation_progress.report("preparing", "inputs"); 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; @@ -2029,6 +2032,7 @@ namespace sd::pipeline { LOG_ERROR("cancelling generation before latent upscale"); return false; } + sd->generation_progress.report("upscaling", "hires"); int64_t upscale_start = ggml_time_ms(); sd::Tensor upscaled_latent; if (sd_version_is_ltxav(sd->version)) { diff --git a/src/runtime/generation_progress.h b/src/runtime/generation_progress.h new file mode 100644 index 0000000000..a211931be7 --- /dev/null +++ b/src/runtime/generation_progress.h @@ -0,0 +1,42 @@ +#ifndef __SD_RUNTIME_GENERATION_PROGRESS_H__ +#define __SD_RUNTIME_GENERATION_PROGRESS_H__ + +#include "stable-diffusion.h" + +namespace sd { + + class GenerationProgress { + sd_generation_progress_cb_t callback_ = nullptr; + void* data_ = nullptr; + int sample_pass_ = 0; + const char* sample_detail_ = ""; + + public: + void set_callback(sd_generation_progress_cb_t callback, void* data) { + callback_ = callback; + data_ = data; + sample_pass_ = 0; + sample_detail_ = ""; + } + + void report(const char* stage, const char* detail = "", int64_t completed = 0, int64_t total = 0, const char* unit = "") const { + if (callback_ != nullptr) { + const sd_generation_progress_t progress{stage, detail, unit, sample_pass_, completed, total}; + callback_(&progress, data_); + } + } + + void begin_sample(int64_t total, const char* detail) { + ++sample_pass_; + sample_detail_ = detail; + report_sample(0, total); + } + + void report_sample(int64_t completed, int64_t total) const { + report("sampling", sample_detail_, completed, total, "steps"); + } + }; + +} // namespace sd + +#endif // __SD_RUNTIME_GENERATION_PROGRESS_H__ diff --git a/src/stable-diffusion.cpp b/src/stable-diffusion.cpp index 436328d77e..4a7af2a2ee 100644 --- a/src/stable-diffusion.cpp +++ b/src/stable-diffusion.cpp @@ -698,6 +698,16 @@ struct sd_ctx_t { StableDiffusionGGML* sd = nullptr; }; +void sd_set_generation_progress_callback(sd_ctx_t* ctx, sd_generation_progress_cb_t cb, void* data) { + if (ctx == nullptr || ctx->sd == nullptr) { + return; + } + std::lock_guard lock(ctx->sd->execution_mutex); + if (!ctx->sd->executing_) { + ctx->sd->generation_progress.set_callback(cb, data); + } +} + sd_ctx_t* new_sd_ctx(const sd_ctx_params_t* sd_ctx_params) { sd_ctx_t* sd_ctx = (sd_ctx_t*)malloc(sizeof(sd_ctx_t)); if (sd_ctx == nullptr) { From f977addf06c1c18f3490bf5ed9b25d2196cd956c Mon Sep 17 00:00:00 2001 From: Wagner Bruna Date: Mon, 9 Nov 2026 11:10:12 -0300 Subject: [PATCH 38/57] fix: disable batch_cfg for Metal when building with upstream ggml (#2311) --- src/pipeline/diffusion_engine.cpp | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index 1ec292c161..e8afbd19a7 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -2440,6 +2440,15 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr denoised = x_t; +#ifdef SD_USE_UPSTREAM_GGML + // workaround for sd.cpp #2149 + const auto diffusion_backends = backend_manager.runtime_backends(SDBackendModule::DIFFUSION); + const bool batch_cfg_backend_ok = std::none_of(diffusion_backends.begin(), diffusion_backends.end(), + [](ggml_backend_t backend) { return sd_backend_is(backend, "MTL"); }); +#else + const bool batch_cfg_backend_ok = true; +#endif + auto denoise = [&](const sd::Tensor& x, float sigma, int step) -> sd::guidance::GuiderOutput { if (get_cancel_flag() == SD_CANCEL_ALL) { LOG_VERBOSE("cancelling generation"); @@ -2762,6 +2771,7 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptrparams.batched_cfg && + batch_cfg_backend_ok && sd_version_is_unet(version) && !uncond.empty() && img_uncond.empty() && From d22dcda1cf439ca42f01a76ddb7c96794dd072ef Mon Sep 17 00:00:00 2001 From: Windyskr <54269955+Windyskr@users.noreply.github.com> Date: Mon, 9 Nov 2026 22:13:35 +0800 Subject: [PATCH 39/57] perf: cache the device type in sd_backend_is_cpu (#2307) --- src/core/ggml_extend_backend.cpp | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/src/core/ggml_extend_backend.cpp b/src/core/ggml_extend_backend.cpp index 47fe80ffd8..70f2ae2ecc 100644 --- a/src/core/ggml_extend_backend.cpp +++ b/src/core/ggml_extend_backend.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include @@ -447,7 +448,19 @@ bool sd_backend_is_cpu(ggml_backend_t backend) { return false; } auto dev = ggml_backend_get_device(backend); - return dev != nullptr && ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU; + if (dev == nullptr) { + return false; + } + // ggml_backend_dev_type() can query the driver on every call (CUDA reads all device + // properties), and this check runs inside per-tensor residency loops. + static std::mutex mutex; + static std::unordered_map is_cpu_by_device; + std::lock_guard lock(mutex); + auto it = is_cpu_by_device.find(dev); + if (it == is_cpu_by_device.end()) { + it = is_cpu_by_device.emplace(dev, ggml_backend_dev_type(dev) == GGML_BACKEND_DEVICE_TYPE_CPU).first; + } + return it->second; } static int cuda_compiled_arch(ggml_backend_t backend) { From cfdf3cf037a9e9ea74f1ce1793e9daf25670baca Mon Sep 17 00:00:00 2001 From: Windyskr <54269955+Windyskr@users.noreply.github.com> Date: Mon, 9 Nov 2026 22:16:04 +0800 Subject: [PATCH 40/57] fix: merge LoRAs in F32 when convert changes the tensor type (#2302) --- src/convert.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/convert.cpp b/src/convert.cpp index b8a62121d5..2e2317ecfe 100644 --- a/src/convert.cpp +++ b/src/convert.cpp @@ -463,11 +463,15 @@ static bool load_model_into_memory(ModelLoader& model_loader, const ggml_type export_type = get_export_tensor_type(model_loader, storage, output_type, type_rules); const bool quant_src = ggml_is_quantized(storage.type); const bool quant_dst = ggml_is_quantized(export_type); + const bool half_src = storage.type == GGML_TYPE_F16 || storage.type == GGML_TYPE_BF16; if (quant_src && quant_dst) { load_type = GGML_TYPE_F16; - } else if (quant_src || quant_dst) { - load_type = quant_src ? export_type : storage.type; - } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32) { + } else if (quant_src) { + load_type = export_type; + } else if (storage.type == GGML_TYPE_F32 || export_type == GGML_TYPE_F32 || + (half_src && export_type != storage.type)) { + // LoRA deltas can be smaller than half an F16/BF16 ulp of the weight. Merge in F32 + // when the export converts the tensor again, so the merged weight is rounded once. load_type = GGML_TYPE_F32; } ggml_tensor* tensor = ggml_new_tensor(ctx, load_type, storage.n_dims, storage.ne); From 008b9a20ca9a3ab1c5d51e8a3456fc0f34435566 Mon Sep 17 00:00:00 2001 From: Windyskr <54269955+Windyskr@users.noreply.github.com> Date: Mon, 9 Nov 2026 22:21:46 +0800 Subject: [PATCH 41/57] fix: apply MiniMax-H3 Diffusers LoRA token_refiner weights (#2306) --- src/name_conversion.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/name_conversion.cpp b/src/name_conversion.cpp index 6bcd483e31..79179f64a9 100644 --- a/src/name_conversion.cpp +++ b/src/name_conversion.cpp @@ -1723,6 +1723,7 @@ std::string convert_tensor_name(std::string name, SDVersion version) { std::vector dit_prefix_vec = { "transformer_blocks", "single_transformer_blocks", + "token_refiner", }; for (const auto& prefix : dit_prefix_vec) { if (starts_with(name, prefix)) { From 991f82558366b35e16ffabf376fa2cba298c12d3 Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 22:33:43 +0800 Subject: [PATCH 42/57] fix: limit default MSVC compilation parallelism --- CMakeLists.txt | 19 +++++++++++++++++-- docs/build.md | 19 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3e210379ae..14901f8f05 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -9,13 +9,28 @@ if (NOT XCODE AND NOT MSVC AND NOT CMAKE_BUILD_TYPE) endif() if (MSVC) + cmake_host_system_information(RESULT SD_LOGICAL_CORES QUERY NUMBER_OF_LOGICAL_CORES) + if (SD_LOGICAL_CORES GREATER 8) + math(EXPR SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES} - 2") + elseif (SD_LOGICAL_CORES GREATER 2) + math(EXPR SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES} - 1") + else() + set(SD_DEFAULT_MSVC_MP_COUNT "${SD_LOGICAL_CORES}") + endif() + if (NOT SD_DEFAULT_MSVC_MP_COUNT GREATER 0) + set(SD_DEFAULT_MSVC_MP_COUNT 1) + endif() + set(SD_MSVC_MP_COUNT "${SD_DEFAULT_MSVC_MP_COUNT}" CACHE STRING "Maximum parallel MSVC compilations per project") + if (NOT SD_MSVC_MP_COUNT MATCHES "^[1-9][0-9]*$" OR SD_MSVC_MP_COUNT GREATER 65536) + message(FATAL_ERROR "SD_MSVC_MP_COUNT must be an integer between 1 and 65536") + endif() add_compile_definitions(_CRT_SECURE_NO_WARNINGS) add_compile_definitions(_SILENCE_CXX17_CODECVT_HEADER_DEPRECATION_WARNING) # /MP is MSVC-only: icx rejects it outright once offloading is enabled. add_compile_options( - $<$,$>:/MP> + $<$,$>:/MP${SD_MSVC_MP_COUNT}> $<$:/utf-8> - $<$,$>:/MP> + $<$,$>:/MP${SD_MSVC_MP_COUNT}> $<$:/utf-8> ) endif() diff --git a/docs/build.md b/docs/build.md index 62966ec6f5..575436a218 100644 --- a/docs/build.md +++ b/docs/build.md @@ -21,6 +21,25 @@ git pull origin master git submodule update --init --recursive ``` +## MSVC compilation parallelism + +MSVC builds limit parallel C/C++ compilation within each project by default. +For `n` logical processors, the limit is `n` when `n <= 2`, `n - 1` when +`3 <= n <= 8`, and `n - 2` when `n > 8`. The default is calculated when the +build directory is first configured with this setting. Override it with +`-DSD_MSVC_MP_COUNT=4`, or pass `-USD_MSVC_MP_COUNT` to recalculate it. + +For Visual Studio builds, keep project-level parallelism at one to avoid +multiplying the per-project limit: + +```shell +cmake --build build --config Release --parallel 1 +``` + +This setting controls MSVC's `/MP` option, not CUDA compiler parallelism or a +hard CPU usage limit. Other build tools, such as Ninja, schedule compilation +jobs separately; use `cmake --build build --parallel N` to limit those jobs. + ## Selecting a GGML source tree By default, sd.cpp builds the patched GGML submodule in `ggml/`. To build with From fbb381da250d47fabb56c3ab5adfc0c2f76ef7e0 Mon Sep 17 00:00:00 2001 From: KIKZILL <121643597+d3Ita@users.noreply.github.com> Date: Mon, 9 Nov 2026 19:34:21 +0400 Subject: [PATCH 43/57] docs: add Prisme to the UIs list (#2274) --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 3a65ee1bfe..46795ab132 100644 --- a/README.md +++ b/README.md @@ -201,6 +201,7 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio - [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel) - [KoboldCpp](https://github.com/LostRuins/koboldcpp) - [Kokon](https://github.com/GabrieleLozupone/kokon) +- [Prisme](https://github.com/d3Ita/prisme) ## Contributors From c44cbec722347adb351a81068f23cfff676d46ce Mon Sep 17 00:00:00 2001 From: 0xConflict <88206582+0xConflict@users.noreply.github.com> Date: Mon, 9 Nov 2026 15:39:31 +0000 Subject: [PATCH 44/57] docs: add Quo to UIs (#2279) --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 46795ab132..015100d829 100644 --- a/README.md +++ b/README.md @@ -201,6 +201,7 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio - [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel) - [KoboldCpp](https://github.com/LostRuins/koboldcpp) - [Kokon](https://github.com/GabrieleLozupone/kokon) +- [Quo](https://github.com/quo-ai/quo) - [Prisme](https://github.com/d3Ita/prisme) ## Contributors From 46da352af0f16a86a1993e1cc938acb4d1e43a34 Mon Sep 17 00:00:00 2001 From: Windyskr <54269955+Windyskr@users.noreply.github.com> Date: Mon, 9 Nov 2026 23:40:26 +0800 Subject: [PATCH 45/57] fix: size layer-split ranges in proportion to device capacity (#2309) --- src/core/layer_split_partition.cpp | 63 +++++++++++++++++++++++++++++- 1 file changed, 61 insertions(+), 2 deletions(-) diff --git a/src/core/layer_split_partition.cpp b/src/core/layer_split_partition.cpp index ff8aaeb602..67b92ee4d8 100644 --- a/src/core/layer_split_partition.cpp +++ b/src/core/layer_split_partition.cpp @@ -73,10 +73,13 @@ namespace sd { return backend_index == 0 ? primary_backend_vram_limit : 0; } + // Capacity of a device that reports no free memory and has no VRAM limit. + static constexpr int64_t graph_cut_layer_split_unknown_capacity = std::numeric_limits::max() / 4; + static std::vector graph_cut_layer_split_backend_capacities(const std::vector& backends, const std::vector& backend_vram_limits, size_t primary_backend_vram_limit) { - std::vector capacities(backends.size(), std::numeric_limits::max() / 4); + std::vector capacities(backends.size(), graph_cut_layer_split_unknown_capacity); constexpr int64_t compute_headroom_bytes = 2ll * 1024 * 1024 * 1024; for (size_t i = 0; i < backends.size(); i++) { ggml_backend_dev_t dev = ggml_backend_get_device(backends[i]); @@ -97,6 +100,57 @@ namespace sd { return capacities; } + // Whether filling the devices in order, each up to `fraction` of its capacity, places every + // segment. Mirrors the placement loop in partition_graph_cut_layer_split(). + static bool graph_cut_layer_split_fits(const std::vector& segment_bytes, + const std::vector& capacities, + double fraction) { + size_t backend = 0; + int64_t used = 0; + for (int64_t bytes : segment_bytes) { + while (backend + 1 < capacities.size() && bytes > 0 && + used + bytes > (int64_t)((double)capacities[backend] * fraction)) { + backend++; + used = 0; + } + if (bytes > 0 && used + bytes > (int64_t)((double)capacities[backend] * fraction)) { + return false; + } + used += bytes; + } + return true; + } + + // Fill limits that size the contiguous ranges in proportion to the device capacities: the same + // fraction of every capacity, as small as possible while all segments still fit. Filling the + // first device up to its capacity would leave only the fixed headroom for compute, which is too + // small for large video graphs. A fraction of 1 is plain capacity-based placement, so every + // layout that fits the capacities is still found. + static std::vector graph_cut_layer_split_fill_limits(const std::vector& segment_bytes, + const std::vector& capacities) { + const bool capacities_known = std::all_of(capacities.begin(), capacities.end(), [](int64_t capacity) { + return capacity < graph_cut_layer_split_unknown_capacity; + }); + if (!capacities_known || !graph_cut_layer_split_fits(segment_bytes, capacities, 1.0)) { + return capacities; + } + double low = 0.0; + double high = 1.0; + for (int i = 0; i < 32; i++) { + const double mid = (low + high) / 2; + if (graph_cut_layer_split_fits(segment_bytes, capacities, mid)) { + high = mid; + } else { + low = mid; + } + } + std::vector limits(capacities.size()); + for (size_t i = 0; i < capacities.size(); i++) { + limits[i] = (int64_t)((double)capacities[i] * high); + } + return limits; + } + bool partition_graph_cut_layer_split(const char* desc, ggml_cgraph* gf, const sd::ggml_graph_cut::Plan& plan, @@ -151,6 +205,11 @@ namespace sd { return param_assignments.count(param) != 0; }); + // Move to the next device at the proportional fill limit; the capacity stays the hard limit. + const std::vector backend_limits = reuse_assignments + ? backend_capacities + : graph_cut_layer_split_fill_limits(segment_param_bytes, backend_capacities); + std::vector backend_by_segment(plan.segments.size(), split_backends[0]); size_t current_backend = 0; int64_t current_used = 0; @@ -158,7 +217,7 @@ namespace sd { int64_t bytes = segment_param_bytes[seg_idx]; while (!reuse_assignments && current_backend + 1 < split_backends.size() && bytes > 0 && - current_used + bytes > backend_capacities[current_backend]) { + current_used + bytes > backend_limits[current_backend]) { current_backend++; current_used = 0; } From d36f72bbe5ec5675da3846e3db966ceb0791d565 Mon Sep 17 00:00:00 2001 From: georgi Date: Mon, 9 Nov 2026 16:41:14 +0100 Subject: [PATCH 46/57] docs: add NodeTool link to README (#2277) --- README.md | 1 + 1 file changed, 1 insertion(+) diff --git a/README.md b/README.md index 015100d829..0ff9339f38 100644 --- a/README.md +++ b/README.md @@ -201,6 +201,7 @@ These projects use `stable-diffusion.cpp` as a backend for their image generatio - [Neural-Pixel](https://github.com/Luiz-Alcantara/Neural-Pixel) - [KoboldCpp](https://github.com/LostRuins/koboldcpp) - [Kokon](https://github.com/GabrieleLozupone/kokon) +- [NodeTool](https://github.com/nodetool-ai/nodetool) - [Quo](https://github.com/quo-ai/quo) - [Prisme](https://github.com/d3Ita/prisme) From 73fef08224c50a49961004dc8c6fee2c0f640bcb Mon Sep 17 00:00:00 2001 From: leejet Date: Mon, 9 Nov 2026 23:41:50 +0800 Subject: [PATCH 47/57] feat: update server web UI with streaming prompt enhancement (#2312) --- examples/server/frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/server/frontend b/examples/server/frontend index d3d83992ff..b07665f893 160000 --- a/examples/server/frontend +++ b/examples/server/frontend @@ -1 +1 @@ -Subproject commit d3d83992ff9bd54c99b11fc36771c13f8637376e +Subproject commit b07665f893047e7b633bbdb69d97c3a9c594076e From a3c49b13be5a64a0b4da2b576af3a17e8ae1ccbf Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 10 Nov 2026 01:58:04 +0800 Subject: [PATCH 48/57] refactor: unify ggml operation helpers around runner context (#2314) --- docs/ggml_extensions.md | 26 +- src/conditioning/conditioner.hpp | 4 +- src/core/ggml_extend.cpp | 703 +++++++++++++++------ src/core/ggml_extend.h | 240 ++++--- src/core/ggml_extensions.cpp | 153 +---- src/core/ggml_extensions.h | 54 +- src/core/ggml_runner.cpp | 16 - src/core/ggml_runner.h | 37 +- src/core/probe/flash_attn.cpp | 6 +- src/model/adapter/lora.hpp | 149 +++-- src/model/adapter/lora_ops.cpp | 77 ++- src/model/adapter/lora_ops.h | 37 +- src/model/adapter/pmid.hpp | 6 +- src/model/audio/wav2vec2.hpp | 10 +- src/model/common/block.hpp | 14 +- src/model/common/ggml_block.hpp | 252 ++------ src/model/common/rope.hpp | 53 +- src/model/common/vsr.hpp | 4 +- src/model/diffusion/anima.hpp | 22 +- src/model/diffusion/animatediff.hpp | 4 +- src/model/diffusion/boogu.hpp | 40 +- src/model/diffusion/control.hpp | 15 +- src/model/diffusion/dit.hpp | 68 +- src/model/diffusion/ernie_image.hpp | 38 +- src/model/diffusion/fibo2.hpp | 20 +- src/model/diffusion/flux.hpp | 48 +- src/model/diffusion/flux3.hpp | 28 +- src/model/diffusion/hidream_o1.hpp | 9 +- src/model/diffusion/hunyuan.hpp | 49 +- src/model/diffusion/ideogram4.hpp | 29 +- src/model/diffusion/iris.hpp | 18 +- src/model/diffusion/krea2.hpp | 30 +- src/model/diffusion/lens.hpp | 18 +- src/model/diffusion/lingbot_video.hpp | 14 +- src/model/diffusion/llada_image.hpp | 22 +- src/model/diffusion/ltxv.hpp | 67 +- src/model/diffusion/ming_image.hpp | 8 +- src/model/diffusion/minimax_h3.hpp | 122 ++-- src/model/diffusion/minit2i.hpp | 24 +- src/model/diffusion/mmdit.hpp | 12 +- src/model/diffusion/moss_image.hpp | 14 +- src/model/diffusion/pid.hpp | 61 +- src/model/diffusion/pixart.hpp | 22 +- src/model/diffusion/qwen_image.hpp | 39 +- src/model/diffusion/qwen_image_2_1.hpp | 40 +- src/model/diffusion/relay_vsr.hpp | 10 +- src/model/diffusion/seedvr2.hpp | 14 +- src/model/diffusion/sefi_image.hpp | 4 +- src/model/diffusion/sensenova_u1.h | 27 +- src/model/diffusion/unet.hpp | 8 +- src/model/diffusion/wan.hpp | 63 +- src/model/diffusion/wan_audio.hpp | 22 +- src/model/diffusion/z_image.hpp | 50 +- src/model/diffusion/z_image_l2p.hpp | 10 +- src/model/te/clip.hpp | 16 +- src/model/te/llada_image_te.hpp | 16 +- src/model/te/llm.hpp | 92 +-- src/model/te/ming_image_te.hpp | 10 +- src/model/te/ph3_image_connector.hpp | 4 +- src/model/te/t5.hpp | 2 +- src/model/upscaler/esrgan.hpp | 4 +- src/model/upscaler/flash_vsr.hpp | 2 +- src/model/upscaler/h3_latent_upscaler.hpp | 14 +- src/model/upscaler/ltx_latent_upscaler.hpp | 34 +- src/model/upscaler/mosr.hpp | 15 +- src/model/vae/auto_encoder_kl.hpp | 14 +- src/model/vae/hunyuan_vae.hpp | 72 +-- src/model/vae/ltx_audio_vae.hpp | 89 +-- src/model/vae/ltx_vae.hpp | 94 ++- src/model/vae/mage_vae.hpp | 92 +-- src/model/vae/minimax_h3_audio_vae.hpp | 12 +- src/model/vae/minimax_h3_vae.hpp | 31 +- src/model/vae/seedvr2_vae.hpp | 32 +- src/model/vae/tae.hpp | 62 +- src/model/vae/wan_vae.hpp | 163 +++-- src/pipeline/diffusion_engine.cpp | 16 +- 76 files changed, 1984 insertions(+), 1832 deletions(-) diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index ea806fb88e..7485938f6a 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -110,7 +110,7 @@ and backend `graph_plan_*` interfaces do not dispatch these extensions. ## Backend services Packed ComfyUI INT4 linear layers use `sd.linear_w4_convrot.v1`, built by -`ggml_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the +`ggml_ext_mul_mat_w4_convrot`. The operator consumes byte-packed I8 weights and the explicit output of `sd.quantize_i8_convrot.v1`; F32 callers get a quantization node during graph construction. Logical output dimensions are stored in the operator attributes, so backend copies of the packed activation do not need to @@ -159,7 +159,15 @@ rebuild a workspace whose actual allocation is smaller than the estimate. ## sd.cpp integration `src/core/ggml_extensions.cpp` and `.h` contain the shared operator schemas, -validation, graph construction helpers, and backend registration entry point. +validation, and backend registration entry point. Graph construction helpers live +in `src/core/ggml_extend.cpp` and `.h`; they take `GGMLRunnerContext*` so backend +selection, compatibility workarounds, and shared ConvRot nodes use the same +context as the calling model. Per-call precision, scaling, attention, and padding +options remain explicit. Model blocks retain weight loading and adapter +orchestration, while the shared helpers own FP8 fallback, embedding lookup, +full/partial RoPE, grouped convolution, and backend-dependent repeat handling. +LoRA/LoKr construction, n-mode products, and Kronecker products stay in +`src/model/adapter/lora_ops.cpp` and `.h`, with the same runner-context interface. CPU kernels and fusion rules live in `src/core/cpu/`, with registration collected in `cpu/registration.cpp`. CUDA-style and Vulkan implementations live in their respective backend directories. @@ -261,13 +269,13 @@ 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 +`ggml_ext_mul_mat_i8_tensorwise`, and `ggml_ext_quantize_i8_convrot`. The latter two constructors are declared and implemented in sd.cpp and create custom nodes. Their former ggml declarations, implementations, ConvRot opcode, and native backend paths have been removed. CPU, CUDA/HIP/MUSA, and Vulkan register kernels and fusion rules through the extension interfaces. -For already quantized activations, `ggml_mul_mat_i8_tensorwise` accepts ordinary +For already quantized activations, `ggml_ext_mul_mat_i8_tensorwise` accepts ordinary I8 weight and activation tensors. Its optional final `input_scale` parameter is a separate F32 tensor with one value per activation row. Weight scales have one value per output channel; weight scales and bias can be null. Existing calls @@ -276,8 +284,8 @@ 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 +`ggml_ext_mul_mat_i8_tensorwise` also accepts F32 activations (quantized into temporary +workspace) or the output of `ggml_ext_quantize_i8_convrot`. ConvRot keeps the existing packed I8 activation rows, padded to a multiple of four rows, followed by F32 row scales. Weight scales remain separate. CPU ConvRot groups are 1, 4, 16, 64, or 256. Unsupported layouts and malformed scale shapes are rejected. Separate CPU quantize @@ -490,8 +498,8 @@ 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 +The `ggml_ext_sage_attn` entry point is declared and implemented in sd.cpp's +graph construction helpers. It constructs `sd.sage_attn.v1` with explicit Q/K/V inputs and scale/mode attributes. ggml has no Sage-specific opcode, allocation rule, or backend dispatch branch. @@ -509,7 +517,7 @@ 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 +The `ggml_ext_sol_attn` entry point lives in the common sd.cpp graph construction files and constructs `sd.sol_attn.v1`. It accepts contiguous F32 Q/K/V tensors with shape `[128, tokens, heads, batch]`, a positive finite scale, and finite tau. The output remains F32 `[128, heads, tokens, batch]`. No current sd.cpp model or diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 7c8b7492f9..1534535990 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -3448,8 +3448,8 @@ struct LTXAVTextProjection : public GGMLBlock { auto video_projection = std::dynamic_pointer_cast(blocks["video_aggregate_embed"]); auto audio_projection = std::dynamic_pointer_cast(blocks["audio_aggregate_embed"]); - auto video_in = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(4096.f / static_cast(kHiddenSize))); - auto audio_in = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(2048.f / static_cast(kHiddenSize))); + auto video_in = ggml_ext_scale(ctx, x, std::sqrt(4096.f / static_cast(kHiddenSize))); + auto audio_in = ggml_ext_scale(ctx, x, std::sqrt(2048.f / static_cast(kHiddenSize))); auto video = video_projection->forward(ctx, video_in); auto audio = audio_projection->forward(ctx, audio_in); return ggml_concat(ctx->ggml_ctx, video, audio, 0); diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 34d35056a9..97561581f1 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -8,51 +8,21 @@ #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 - // swap 0th and nth axis - a = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0)); - int64_t ne1 = a->ne[1]; - int64_t ne2 = a->ne[2]; - int64_t ne3 = a->ne[3]; - // make 2D - a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1))); - - ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b))); - - // reshape output (same shape as a after permutation except first dim) - result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3); - // swap back 0th and nth axis - result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0); - return result; -} - -ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b) { - return ggml_mul(ctx, - ggml_interpolate(ctx, - a, - a->ne[0] * b->ne[0], - a->ne[1] * b->ne[1], - a->ne[2] * b->ne[2], - a->ne[3] * b->ne[3], - GGML_SCALE_MODE_NEAREST), - b); -} - -ggml_tensor* ggml_ext_cont(ggml_context* ctx, - ggml_tensor* x) { +ggml_tensor* ggml_ext_cont(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + auto ctx = runner_ctx->ggml_ctx; if (ggml_is_contiguous(x)) { return x; } return ggml_cont(ctx, x); } -ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx, +ggml_tensor* ggml_ext_torch_permute(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int axis0, int axis1, int axis2, int axis3) { + auto ctx = runner_ctx->ggml_ctx; int torch_axes[4] = {axis0, axis1, axis2, axis3}; int ggml_axes[4] = {0}; @@ -71,12 +41,13 @@ ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx, return ggml_permute(ctx, x, ggml_axes[0], ggml_axes[1], ggml_axes[2], ggml_axes[3]); } -ggml_tensor* ggml_ext_slice(ggml_context* ctx, +ggml_tensor* ggml_ext_slice(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int dim, int64_t start, int64_t end, bool cont) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(dim >= 0 && dim < 4); if (x->ne[dim] == 1) { return x; @@ -106,11 +77,12 @@ ggml_tensor* ggml_ext_slice(ggml_context* ctx, return x; } -std::vector ggml_ext_chunk(ggml_context* ctx, +std::vector ggml_ext_chunk(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int num, int64_t dim, bool cont) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(dim >= 0 && dim < 4); GGML_ASSERT(x->ne[dim] % num == 0); @@ -133,7 +105,8 @@ std::vector ggml_ext_chunk(ggml_context* ctx, return chunks; } -ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first) { +ggml_tensor* ggml_ext_silu_act(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool gate_first) { + auto ctx = runner_ctx->ggml_ctx; // x: [ne3, ne2, ne1, ne0] // return: [ne3, ne2, ne1, ne0/2] @@ -144,31 +117,25 @@ ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_firs return gate_first ? ggml_swiglu(ctx, x) : ggml_swiglu_swapped(ctx, x); } -ggml_tensor* ggml_ext_situ(GGMLRunnerContext* ctx, ggml_tensor* hidden, ggml_tensor* gate) { +ggml_tensor* ggml_ext_situ(GGMLRunnerContext* runner_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)) { + if (runner_ctx->backend != nullptr) { + auto out = gate ? ggml_glu_split(runner_ctx->ggml_ctx, gate, hidden, GGML_GLU_OP_SITU) + : ggml_glu(runner_ctx->ggml_ctx, hidden, GGML_GLU_OP_SITU, true); + if (ggml_backend_supports_op(runner_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); + gate = ggml_ext_slice(runner_ctx, hidden, 0, dim, 2 * dim); + hidden = ggml_ext_slice(runner_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, - ggml_tensor* a) { - const float eps = 1e-6f; // default eps parameter - return ggml_group_norm(ctx, a, 32, eps); + auto a = ggml_scale(runner_ctx->ggml_ctx, ggml_tanh(runner_ctx->ggml_ctx, ggml_scale(runner_ctx->ggml_ctx, gate, 0.25f)), 4.f); + a = ggml_mul(runner_ctx->ggml_ctx, a, ggml_sigmoid(runner_ctx->ggml_ctx, gate)); + auto b = ggml_scale(runner_ctx->ggml_ctx, ggml_tanh(runner_ctx->ggml_ctx, ggml_scale(runner_ctx->ggml_ctx, hidden, 1.f / 25.f)), 25.f); + return ggml_mul(runner_ctx->ggml_ctx, a, b); } static bool ggml_ext_is_padded_1d(const ggml_tensor* x) { @@ -177,30 +144,28 @@ 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* ggml_ext_modulate(GGMLRunnerContext* runner_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"))) { + (sd_backend_is_cpu(runner_ctx->backend) || sd_backend_is(runner_ctx->backend, "CUDA") || + sd_backend_is(runner_ctx->backend, "ROCm") || sd_backend_is(runner_ctx->backend, "MUSA") || sd_backend_is(runner_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); + return ggml_ext_op_create(runner_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)); + x = ggml_add(runner_ctx->ggml_ctx, x, ggml_mul(runner_ctx->ggml_ctx, x, scale)); if (shift != nullptr) { - x = ggml_add(ctx->ggml_ctx, x, shift); + x = ggml_add(runner_ctx->ggml_ctx, x, shift); } return x; } -ggml_tensor* ggml_ext_scale(ggml_context* ctx, - ggml_tensor* x, - float factor, - bool inplace) { +ggml_tensor* ggml_ext_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* x, float factor, bool inplace) { + auto ctx = runner_ctx->ggml_ctx; if (!ggml_ext_is_padded_1d(x)) { x = ggml_cont(ctx, x); } @@ -212,9 +177,8 @@ ggml_tensor* ggml_ext_scale(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_gelu(ggml_context* ctx, - ggml_tensor* x, - bool inplace) { +ggml_tensor* ggml_ext_gelu(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace) { + auto ctx = runner_ctx->ggml_ctx; if (!ggml_is_contiguous(x)) { x = ggml_cont(ctx, x); } @@ -226,9 +190,8 @@ ggml_tensor* ggml_ext_gelu(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, - ggml_tensor* x, - bool inplace) { +ggml_tensor* ggml_ext_gelu_quick(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace) { + auto ctx = runner_ctx->ggml_ctx; if (!ggml_is_contiguous(x)) { x = ggml_cont(ctx, x); } @@ -240,14 +203,16 @@ ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_linear(ggml_context* ctx, +ggml_tensor* ggml_ext_linear(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, bool force_prec_f32, float scale) { + auto ctx = runner_ctx->ggml_ctx; + w = ggml_ext_prepare_linear_weight(runner_ctx, w, x, force_prec_f32); if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); + x = ggml_ext_scale(runner_ctx, x, scale); } if (x->ne[2] * x->ne[3] > 1024) { // workaround: avoid ggml cuda error @@ -266,7 +231,7 @@ ggml_tensor* ggml_ext_linear(ggml_context* ctx, } } if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); + x = ggml_ext_scale(runner_ctx, x, 1.f / scale); } if (b != nullptr) { x = ggml_add_inplace(ctx, x, b); @@ -274,25 +239,30 @@ ggml_tensor* ggml_ext_linear(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, +ggml_tensor* ggml_ext_linear_i8_tensorwise(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* weight_scale, ggml_tensor* b, int convrot_group_size, float scale) { + auto ctx = runner_ctx->ggml_ctx; #ifndef SD_USE_UPSTREAM_GGML + x = ggml_ext_prepare_i8_input(runner_ctx, x); + if (x->type == GGML_TYPE_F32 && convrot_group_size > 0 && scale == 1.f) { + x = ggml_ext_quantize_i8_convrot(runner_ctx, x, convrot_group_size); + } GGML_ASSERT(x->type == GGML_TYPE_F32 || (x->type == GGML_TYPE_I8 && scale == 1.f)); if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); + x = ggml_ext_scale(runner_ctx, x, scale); } ggml_tensor* fused_bias = scale == 1.f ? b : nullptr; auto mul_mat = [&](ggml_tensor* input) { if (input->type == GGML_TYPE_F32 && convrot_group_size > 0) { - input = ggml_quantize_i8_convrot(ctx, input, convrot_group_size); + input = ggml_ext_quantize_i8_convrot(runner_ctx, input, convrot_group_size); } - return ggml_mul_mat_i8_tensorwise(ctx, w, input, weight_scale, fused_bias, convrot_group_size); + return ggml_ext_mul_mat_i8_tensorwise(runner_ctx, w, input, weight_scale, fused_bias, convrot_group_size); }; if (x->ne[2] * x->ne[3] > 1024) { int64_t ne2 = x->ne[2]; @@ -305,7 +275,7 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, } if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); + x = ggml_ext_scale(runner_ctx, x, 1.f / scale); if (b != nullptr) { x = ggml_add_inplace(ctx, x, b); } @@ -323,8 +293,7 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, #endif } -ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, - ggml_backend_t backend, +ggml_tensor* ggml_ext_pad_ext(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int lp0, int rp0, @@ -336,6 +305,8 @@ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, int rp3, bool circular_x, bool circular_y) { + auto ctx = runner_ctx->ggml_ctx; + auto backend = runner_ctx->backend; if (circular_x && circular_y) { return ggml_pad_ext_circular(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); } @@ -365,7 +336,7 @@ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_pad(ggml_context* ctx, +ggml_tensor* ggml_ext_pad(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int p0, int p1, @@ -373,7 +344,7 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx, int p3, bool circular_x, bool circular_y) { - return ggml_ext_pad_ext(ctx, nullptr, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); + return ggml_ext_pad_ext(runner_ctx, x, 0, p0, 0, p1, 0, p2, 0, p3, circular_x, circular_y); } static ggml_tensor* conv_1d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, int s0, int p0, int d0, bool force_prec_f32) { @@ -388,7 +359,7 @@ static ggml_tensor* conv_1d(ggml_context* ctx, ggml_tensor* x, ggml_tensor* w, i return ggml_cont(ctx, ggml_permute(ctx, result, 0, 2, 1, 3)); } -ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, +ggml_tensor* ggml_ext_conv_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -397,6 +368,7 @@ ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, int d0, int64_t groups, bool force_prec_f32) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(s0 > 0 && p0 >= 0 && d0 > 0 && groups > 0); GGML_ASSERT(x->type == GGML_TYPE_F32 && x->ne[3] == 1 && w->ne[3] == 1); GGML_ASSERT(x->ne[1] % groups == 0 && w->ne[2] % groups == 0); @@ -427,7 +399,7 @@ ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, ggml_tensor* w_i = ggml_view_3d(ctx, w, w->ne[0], ic_g, oc_g, w->nb[1], w->nb[2], group * oc_g * w->nb[2]); outputs.push_back(conv_1d(ctx, x_i, w_i, s0, p0, d0, force_prec_f32)); } - result = ggml_ext_vec_concat(ctx, outputs, 1); + result = ggml_ext_vec_concat(runner_ctx, outputs, 1); } if (b != nullptr) { if (!ggml_is_contiguous(b)) { @@ -439,7 +411,7 @@ ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, return result; } -ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, +ggml_tensor* ggml_ext_conv_2d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -453,15 +425,16 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, bool circular_x, bool circular_y, float scale) { + auto ctx = runner_ctx->ggml_ctx; if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); + x = ggml_ext_scale(runner_ctx, x, scale); } if (w->ne[2] != x->ne[2] && ggml_n_dims(w) == 2) { w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], w->ne[1]); } if ((p0 != 0 || p1 != 0) && (circular_x || circular_y)) { - x = ggml_ext_pad_ext(ctx, nullptr, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); + x = ggml_ext_pad_ext(runner_ctx, x, p0, p0, p1, p1, 0, 0, 0, 0, circular_x, circular_y); p0 = 0; p1 = 0; } @@ -472,7 +445,7 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, x = ggml_conv_2d(ctx, w, x, s0, s1, p0, p1, d0, d1); } if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); + x = ggml_ext_scale(runner_ctx, x, 1.f / scale); } if (b != nullptr) { b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); @@ -481,8 +454,7 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, - ggml_backend_t backend, +ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -498,9 +470,18 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d2, bool force_prec_f32, bool direct, - float scale) { + float scale, + bool allow_conv2d) { + auto ctx = runner_ctx->ggml_ctx; + auto backend = runner_ctx->backend; + if (allow_conv2d && w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == IC && p2 == 0) { + auto x2 = ggml_reshape_4d(ctx, ggml_ext_cont(runner_ctx, x), x->ne[0], x->ne[1], IC, 1); + auto w2 = ggml_reshape_4d(ctx, w, w->ne[0], w->ne[1], IC, w->ne[3] / IC); + x2 = ggml_ext_conv_2d(runner_ctx, x2, w2, b, s0, s1, p0, p1, d0, d1, runner_ctx->conv2d_direct_enabled, false, false, scale); + return ggml_reshape_4d(ctx, x2, x2->ne[0], x2->ne[1], 1, w->ne[3] / IC); + } if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, scale); + x = ggml_ext_scale(runner_ctx, x, scale); } if (direct) { int64_t OC = w->ne[3] / IC; @@ -551,7 +532,7 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, } if (scale != 1.f) { - x = ggml_ext_scale(ctx, x, 1.f / scale); + x = ggml_ext_scale(runner_ctx, x, 1.f / scale); } if (b != nullptr) { b = ggml_reshape_4d(ctx, b, 1, 1, 1, b->ne[0]); // [OC, 1, 1, 1] @@ -560,25 +541,10 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - int s2, - int p2, - int d2) { - x = ggml_conv_2d(ctx, w, x, 1, s2, 0, p2, 1, d2); // [N, OC, T, OH * OW] - if (b != nullptr) { - b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); - x = ggml_add(ctx, x, b); - } - return x; // [N, OC, T, OH * OW] -} - -std::vector split_qkv(ggml_context* ctx, - ggml_tensor* qkv) { - qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]); // [N, L, 3, C] - qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // [3, N, L, C] +std::vector ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) { + auto ctx = runner_ctx->ggml_ctx; + qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]); // [N, L, 3, C] + qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // [3, N, L, C] int64_t offset = qkv->nb[2] * qkv->ne[2]; auto q = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0); // [N, L, C] @@ -587,16 +553,16 @@ std::vector split_qkv(ggml_context* ctx, return {q, k, v}; } -std::vector split_image_qkv(ggml_context* ctx, - ggml_tensor* qkv) { +std::vector ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) { + auto ctx = runner_ctx->ggml_ctx; int64_t W = qkv->ne[0]; int64_t H = qkv->ne[1]; int64_t C = qkv->ne[2] / 3; int64_t N = qkv->ne[3]; int64_t nb1 = qkv->nb[1]; int64_t nb2 = qkv->nb[2]; - qkv = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N); // [N, 3, C, H*W] - qkv = ggml_cont(ctx, ggml_ext_torch_permute(ctx, qkv, 0, 1, 3, 2)); // [3, N, C, H*W] + qkv = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N); // [N, 3, C, H*W] + qkv = ggml_cont(ctx, ggml_ext_torch_permute(runner_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[2], offset * 0); // [N, C, H, W] @@ -605,45 +571,38 @@ std::vector split_image_qkv(ggml_context* ctx, return {q, k, v}; } -ggml_tensor* ggml_ext_full(ggml_context* ctx, +ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx, float value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + auto ctx = runner_ctx->ggml_ctx; auto one = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:one"); - auto t = ggml_ext_scale(ctx, one, value); // [1,] + auto t = ggml_ext_scale(runner_ctx, one, value); // [1,] t = ggml_repeat_4d(ctx, t, ne0, ne1, ne2, ne3); // [ne0, ne1, ne2, ne3] return t; } -ggml_tensor* ggml_ext_zeros(ggml_context* ctx, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - return ggml_ext_full(ctx, 0.f, ne0, ne1, ne2, ne3); +ggml_tensor* ggml_ext_zeros(GGMLRunnerContext* runner_ctx, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return ggml_ext_full(runner_ctx, 0.f, ne0, ne1, ne2, ne3); } -ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx, - ggml_tensor* x) { - return ggml_ext_zeros(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); +ggml_tensor* ggml_ext_zeros_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + return ggml_ext_zeros(runner_ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); } -ggml_tensor* ggml_ext_ones(ggml_context* ctx, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - return ggml_ext_full(ctx, 1.f, ne0, ne1, ne2, ne3); +ggml_tensor* ggml_ext_ones(GGMLRunnerContext* runner_ctx, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + return ggml_ext_full(runner_ctx, 1.f, ne0, ne1, ne2, ne3); } -ggml_tensor* ggml_ext_ones_like(ggml_context* ctx, - ggml_tensor* x) { - return ggml_ext_ones(ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); +ggml_tensor* ggml_ext_ones_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + return ggml_ext_ones(runner_ctx, x->ne[0], x->ne[1], x->ne[2], x->ne[3]); } -ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a) { +ggml_tensor* ggml_ext_cast_f32(GGMLRunnerContext* runner_ctx, ggml_tensor* a) { + auto ctx = runner_ctx->ggml_ctx; + auto backend = runner_ctx->backend; if (sd_backend_is(backend, "Vulkan")) { auto zero_index = ggml_get_tensor(ctx, "ggml_runner_build_in_tensor:zero_int"); auto out = ggml_reshape_1d(ctx, a, ggml_nelements(a)); @@ -653,7 +612,7 @@ ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_t return out; } else { auto out = ggml_reshape_2d(ctx, a, 1, ggml_nelements(a)); - ggml_tensor* one = ggml_ext_ones(ctx, 1, 1, 1, 1); // [1,] + ggml_tensor* one = ggml_ext_ones(runner_ctx, 1, 1, 1, 1); // [1,] if (ggml_is_transposed(out)) { out = ggml_mul_mat(ctx, one, out); } else { @@ -664,8 +623,7 @@ ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_t } } -ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, - ggml_backend_t backend, +ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* runner_ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, @@ -674,8 +632,13 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, bool skip_reshape, bool flash_attn, float kv_scale, - bool sage_attn, - bool* used_flash_attn) { // avoid overflow + bool* used_flash_attn) { + auto ctx = runner_ctx->ggml_ctx; + auto backend = runner_ctx->backend; + if (runner_ctx->attn_scale > 0.f) { + kv_scale = runner_ctx->attn_scale; + } + const bool sage_attn = runner_ctx->sage_attn_enabled; if (used_flash_attn != nullptr) { *used_flash_attn = false; } @@ -693,13 +656,13 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, d_head = C / n_head; n_kv_head = k->ne[0] / d_head; - q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N); // [N, L_q, n_head, d_head] - q = ggml_ext_cont(ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, L_q, d_head] - q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N); // [N * n_head, L_q, d_head] + q = ggml_reshape_4d(ctx, q, d_head, n_head, L_q, N); // [N, L_q, n_head, d_head] + q = ggml_ext_cont(runner_ctx, ggml_permute(ctx, q, 0, 2, 1, 3)); // [N, n_head, L_q, d_head] + q = ggml_reshape_3d(ctx, q, d_head, L_q, n_head * N); // [N * n_head, L_q, d_head] - k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] - k = ggml_ext_cont(ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_kv_head, L_k, d_head] - k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N); // [N * n_kv_head, L_k, d_head] + k = ggml_reshape_4d(ctx, k, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] + k = ggml_ext_cont(runner_ctx, ggml_permute(ctx, k, 0, 2, 1, 3)); // [N, n_kv_head, L_k, d_head] + k = ggml_reshape_3d(ctx, k, d_head, L_k, n_kv_head * N); // [N * n_kv_head, L_k, d_head] v = ggml_reshape_4d(ctx, v, d_head, n_kv_head, L_k, N); // [N, L_k, n_kv_head, d_head] } else { @@ -726,17 +689,17 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, 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_ext_scale(runner_ctx, k_in, kv_scale); } k_in = ggml_cast(ctx, k_in, GGML_TYPE_F16); if (kv_scale != 1.0f) { - v_in = ggml_ext_scale(ctx, v_in, kv_scale); + v_in = ggml_ext_scale(runner_ctx, v_in, kv_scale); } 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_f16 = ggml_cast(ctx, ggml_ext_cont(runner_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); @@ -763,19 +726,19 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, } ggml_prec_set_acc(out, GGML_PREC_F32); if (kv_scale != 1.0f) { - out = ggml_ext_scale(ctx, out, 1.0f / kv_scale); + out = ggml_ext_scale(runner_ctx, out, 1.0f / kv_scale); } if (pad_head) { - out = ggml_ext_slice(ctx, out, 0, 0, d_head); + out = ggml_ext_slice(runner_ctx, out, 0, 0, d_head); } return out; }; #ifndef SD_USE_UPSTREAM_GGML if (sage_attn && mask == nullptr && d_head > 0 && d_head <= 128) { - auto q_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, q->type == GGML_TYPE_F32 ? q : ggml_cast(ctx, q, GGML_TYPE_F32)), d_head, L_q, n_head, N); - auto k_in = ggml_reshape_4d(ctx, ggml_ext_cont(ctx, k->type == GGML_TYPE_F32 ? k : ggml_cast(ctx, k, GGML_TYPE_F32)), d_head, L_k, n_kv_head, N); - auto v_in = ggml_ext_cont(ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); + auto q_in = ggml_reshape_4d(ctx, ggml_ext_cont(runner_ctx, q->type == GGML_TYPE_F32 ? q : ggml_cast(ctx, q, GGML_TYPE_F32)), d_head, L_q, n_head, N); + auto k_in = ggml_reshape_4d(ctx, ggml_ext_cont(runner_ctx, k->type == GGML_TYPE_F32 ? k : ggml_cast(ctx, k, GGML_TYPE_F32)), d_head, L_k, n_kv_head, N); + auto v_in = ggml_ext_cont(runner_ctx, ggml_permute(ctx, v, 0, 2, 1, 3)); const int64_t padded_head = d_head <= 64 ? 64 : 128; if ((padded_head != d_head || kv_scale != 1.0f) && v_in->type != GGML_TYPE_F32) { v_in = ggml_cast(ctx, v_in, GGML_TYPE_F32); @@ -787,15 +750,15 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, v_in = ggml_pad(ctx, v_in, static_cast(padded_head - d_head), 0, 0, 0); } if (kv_scale != 1.0f) { - k_in = ggml_ext_scale(ctx, k_in, kv_scale); - v_in = ggml_ext_scale(ctx, v_in, kv_scale); + k_in = ggml_ext_scale(runner_ctx, k_in, kv_scale); + v_in = ggml_ext_scale(runner_ctx, v_in, kv_scale); } v_in = ggml_cast(ctx, v_in, GGML_TYPE_F16); - auto out = ggml_sage_attn(ctx, q_in, k_in, v_in, scale / kv_scale, GGML_SAGE_ATTN_AUTO); + auto out = ggml_ext_sage_attn(runner_ctx, q_in, k_in, v_in, scale / kv_scale, GGML_SAGE_ATTN_AUTO); if (ggml_backend_supports_op(backend, out)) { - kqv = kv_scale != 1.0f ? ggml_ext_scale(ctx, out, 1.0f / kv_scale) : out; + kqv = kv_scale != 1.0f ? ggml_ext_scale(runner_ctx, out, 1.0f / kv_scale) : out; if (padded_head != d_head) { - kqv = ggml_ext_slice(ctx, kqv, 0, 0, d_head); + kqv = ggml_ext_slice(runner_ctx, kqv, 0, 0, d_head); } } } @@ -833,8 +796,8 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, // if (flash_attn) { // LOG_VERBOSE("fallback to default attention, L_q:%d L_k:%d n_head:%d C:%d d_head:%d N:%d", L_q, L_k, n_head, C, d_head, N); // } - v = ggml_ext_cont(ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_kv_head, d_head, L_k] - v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N); // [N * n_kv_head, d_head, L_k] + v = ggml_ext_cont(runner_ctx, ggml_permute(ctx, v, 1, 2, 0, 3)); // [N, n_kv_head, d_head, L_k] + v = ggml_reshape_3d(ctx, v, L_k, d_head, n_kv_head * N); // [N * n_kv_head, d_head, L_k] auto kq = ggml_mul_mat(ctx, k, q); // [N * n_head, L_q, L_k] ggml_prec_set_acc(kq, GGML_PREC_F32); @@ -850,18 +813,19 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, kqv = ggml_permute(ctx, kqv, 0, 2, 1, 3); // [N, L_q, n_head, d_head] } - kqv = ggml_ext_cont(ctx, kqv); + kqv = ggml_ext_cont(runner_ctx, kqv); kqv = ggml_reshape_3d(ctx, kqv, d_head * n_head, L_q, N); // [N, L_q, C] return kqv; } -ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, +ggml_tensor* ggml_ext_layer_norm(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, float eps) { - x = ggml_norm(ctx, x, eps); + auto ctx = runner_ctx->ggml_ctx; + x = ggml_norm(ctx, x, eps); if (w != nullptr) { x = ggml_mul_inplace(ctx, x, w); if (b != nullptr) { @@ -871,12 +835,13 @@ ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, +ggml_tensor* ggml_ext_group_norm(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, int num_groups, float eps) { + auto ctx = runner_ctx->ggml_ctx; if (ggml_n_dims(x) >= 3 && w != nullptr && b != nullptr) { w = ggml_reshape_4d(ctx, w, 1, 1, w->ne[0], 1); b = ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1); @@ -890,19 +855,18 @@ ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, return x; } -ggml_tensor* ggml_ext_timestep_embedding( - ggml_context* ctx, - ggml_tensor* timesteps, - int dim, - int max_period, - float time_factor) { - timesteps = ggml_ext_scale(ctx, timesteps, time_factor); +ggml_tensor* ggml_ext_timestep_embedding(GGMLRunnerContext* runner_ctx, + ggml_tensor* timesteps, + int dim, + int max_period, + float time_factor) { + auto ctx = runner_ctx->ggml_ctx; + timesteps = ggml_ext_scale(runner_ctx, timesteps, time_factor); return ggml_timestep_embedding(ctx, timesteps, dim, max_period); } -ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx, - std::vector& tensors, - int dim) { +ggml_tensor* ggml_ext_vec_concat(GGMLRunnerContext* runner_ctx, std::vector& tensors, int dim) { + auto ctx = runner_ctx->ggml_ctx; while (tensors.size() > 1) { std::vector next_level; for (size_t i = 0; i < tensors.size(); i += 2) { @@ -916,3 +880,388 @@ ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx, } return tensors[0]; } + +#ifndef SD_USE_UPSTREAM_GGML +ggml_tensor* ggml_ext_mul_mat_i8_tensorwise(GGMLRunnerContext* runner_ctx, + ggml_tensor* w, + ggml_tensor* x, + ggml_tensor* ws, + ggml_tensor* bias, + int group, + ggml_tensor* xs) { + auto ctx = runner_ctx->ggml_ctx; + 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]; + } + SDI8Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? SD_I8_PACKED : 0}; + ggml_tensor* inputs[5] = {w, x}; + int count = 2; + if (xs) { + attr.flags |= SD_I8_INPUT_SCALE; + inputs[count++] = xs; + } + if (ws) { + attr.flags |= SD_I8_WEIGHT_SCALE; + inputs[count++] = ws; + } + if (bias) { + attr.flags |= SD_I8_BIAS; + inputs[count++] = bias; + } + return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr); +} + +ggml_tensor* ggml_ext_quantize_i8_convrot(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int group) { + auto ctx = runner_ctx->ggml_ctx; + const auto key = std::make_pair(x, group); + auto cached = runner_ctx->int8_convrot_cache.find(key); + if (cached != runner_ctx->int8_convrot_cache.end()) { + return cached->second; + } + int32_t attr = group; + auto result = ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr); + if (result) { + runner_ctx->int8_convrot_cache.emplace(key, result); + } + return result; +} + +ggml_tensor* ggml_ext_mul_mat_w4_convrot(GGMLRunnerContext* runner_ctx, + ggml_tensor* w, + ggml_tensor* x, + ggml_tensor* ws, + ggml_tensor* sc, + ggml_tensor* sr, + ggml_tensor* bias, + SDW4Kind kind, + int group) { + auto ctx = runner_ctx->ggml_ctx; + if (!w || !x || group < 0 || (group && !sd_ggml_valid_convrot_group(group))) + return nullptr; + const int quant_group = group ? group : 1; + if (x->type == GGML_TYPE_F32) + x = ggml_ext_quantize_i8_convrot(runner_ctx, x, quant_group); + int packed_group; + if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group) + return nullptr; + const auto* logical = x->src[0]; + SDW4Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0}; + ggml_tensor* inputs[6] = {w, x, ws}; + int count = 3; + if (kind == SDW4Kind::W4A8) { + inputs[count++] = sc; + inputs[count++] = sr; + } else if (sc || sr) { + return nullptr; + } + if (bias) + inputs[count++] = bias; + return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &attr); +} + +ggml_tensor* ggml_ext_sage_attn(GGMLRunnerContext* runner_ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + float scale, + ggml_sage_attn_mode mode) { + auto ctx = runner_ctx->ggml_ctx; + 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); +} + +ggml_tensor* ggml_ext_sol_attn(GGMLRunnerContext* runner_ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + float scale, + float tau) { + auto ctx = runner_ctx->ggml_ctx; + 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); +} + +ggml_tensor* ggml_ext_window_attn(GGMLRunnerContext* runner_ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + const SDWindowAttentionParams& params) { + auto ctx = runner_ctx->ggml_ctx; + auto backend = runner_ctx->backend; + ggml_tensor* inputs[] = {q, k, v}; + ggml_ext_tensor_desc out; + if (!backend || !sd_ggml_window_attn_infer(inputs, 3, ¶ms, &out)) + return nullptr; + sd_ggml_register_extensions(backend); + auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, ¶ms); + return ggml_backend_supports_op(backend, result) ? result : nullptr; +} + +#endif + +ggml_tensor* ggml_ext_rope_apply(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* pe, + bool rope_interleaved) { + auto ctx = runner_ctx->ggml_ctx; + const int64_t rot_dim = pe->ne[2] * 2; + GGML_ASSERT(rot_dim > 0 && rot_dim <= x->ne[0]); +#ifndef SD_USE_UPSTREAM_GGML + auto backend = runner_ctx->backend; + 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]) { + sd_ggml_register_extensions(backend); + ggml_tensor* inputs[] = {x, pe}; + int32_t interleaved = rope_interleaved ? 1 : 0; + auto out = ggml_ext_op_create(ctx, sd_ggml_rope_op(), inputs, 2, &interleaved); + if (out && ggml_backend_supports_op(backend, out)) { + return out; + } + } +#endif + if (rot_dim < x->ne[0]) { + auto rotated = ggml_ext_rope_apply(runner_ctx, ggml_ext_slice(runner_ctx, x, 0, 0, rot_dim), pe, rope_interleaved); + auto tail = ggml_ext_slice(runner_ctx, x, 0, rot_dim, x->ne[0]); + tail = ggml_cont(ctx, ggml_permute(ctx, tail, 0, 2, 1, 3)); + tail = ggml_reshape_3d(ctx, tail, tail->ne[0], tail->ne[1], tail->ne[2] * tail->ne[3]); + return ggml_concat(ctx, rotated, tail, 0); + } + // x: [N, L, n_head, d_head] + // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]] + int64_t d_head = x->ne[0]; + int64_t n_head = x->ne[1]; + int64_t L = x->ne[2]; + int64_t N = x->ne[3]; + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N, n_head, L, d_head] + if (rope_interleaved) { + x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N); // [N * n_head, L, d_head/2, 2] + x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2)); // [2, N * n_head, L, d_head/2] + } else { + auto x_0 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], 0); + auto x_1 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], (d_head / 2) * x->nb[0]); + + auto pe_2 = ggml_cont(ctx, ggml_permute(ctx, pe, 2, 3, 0, 1)); + auto cos = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], 0); + auto sin = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], pe_2->nb[3]); + + auto part_0 = ggml_sub(ctx, ggml_mul(ctx, x_0, cos), ggml_mul(ctx, x_1, sin)); + auto part_1 = ggml_add(ctx, ggml_mul(ctx, x_0, sin), ggml_mul(ctx, x_1, cos)); + return ggml_concat(ctx, part_0, part_1, 0); + } + + int64_t offset = x->nb[2] * x->ne[2]; + auto x_0 = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0); // [N * n_head, L, d_head/2] + auto x_1 = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1); // [N * n_head, L, d_head/2] + x_0 = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]); // [N * n_head, L, d_head/2, 1] + x_1 = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]); // [N * n_head, L, d_head/2, 1] + auto temp_x = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]); + x_0 = ggml_repeat(ctx, x_0, temp_x); // [N * n_head, L, d_head/2, 2] + x_1 = ggml_repeat(ctx, x_1, temp_x); // [N * n_head, L, d_head/2, 2] + + pe = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2)); // [2, L, d_head/2, 2] + offset = pe->nb[2] * pe->ne[2]; + auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0); // [L, d_head/2, 2] + auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1); // [L, d_head/2, 2] + + auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1)); // [N * n_head, L, d_head/2, 2] + x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N); // [N*n_head, L, d_head] + return x_out; +} + +ggml_tensor* ggml_ext_repeat_4d(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + auto ctx = runner_ctx->ggml_ctx; + if ((x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && sd_backend_is(runner_ctx->backend, "Vulkan")) { + auto repeated = ggml_repeat_4d(ctx, ggml_cast(ctx, x, GGML_TYPE_F32), ne0, ne1, ne2, ne3); + return ggml_cast(ctx, repeated, x->type); + } + return ggml_repeat_4d(ctx, x, ne0, ne1, ne2, ne3); +} + +ggml_tensor* ggml_ext_repeat(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* shape) { + return ggml_ext_repeat_4d(runner_ctx, x, shape->ne[0], shape->ne[1], shape->ne[2], shape->ne[3]); +} + +bool ggml_ext_force_prec_f32(GGMLRunnerContext* runner_ctx, + bool requested, + uint32_t flags) { + return requested || + ((flags & GGML_PRECISION_F32_VULKAN) != 0 && sd_backend_is(runner_ctx->backend, "Vulkan")) || + ((flags & GGML_PRECISION_F32_ROCM) != 0 && sd_backend_is(runner_ctx->backend, "ROCm")); +} + +static void ggml_ext_get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) { + const int group_size = static_cast(reinterpret_cast(userdata)); + const auto* weight = dst->src[0]; + const auto* input_ids = static_cast(dst->src[1]->data); + const auto* scales = static_cast(dst->src[2]->data); + const bool scalar_scale = ggml_nelements(dst->src[2]) == 1; + const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast(group_size)) : 1.f; + for (int64_t row = ith; row < dst->ne[1]; row += nth) { + const int32_t token = input_ids[row]; + GGML_ASSERT(token >= 0 && token < weight->ne[1]); + const auto* src = reinterpret_cast(static_cast(weight->data) + token * weight->nb[1]); + auto* out = reinterpret_cast(static_cast(dst->data) + row * dst->nb[1]); + const float scale = scales[scalar_scale ? 0 : token] * normalization; + for (int64_t i = 0; i < weight->ne[0]; ++i) { + out[i] = static_cast(src[i]) * scale; + } + // The regular Hadamard rotation is symmetric and its own inverse. + for (int stride = 1; stride < group_size; stride *= 4) { + for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) { + for (int j = 0; j < stride; ++j) { + float* values = out + base + j; + const float a = values[0]; + const float b = values[stride]; + const float c = values[2 * stride]; + const float d = values[3 * stride]; + values[0] = a + b + c - d; + values[stride] = a + b - c + d; + values[2 * stride] = a - b + c + d; + values[3 * stride] = -a + b + c + d; + } + } + } + } +} + +ggml_tensor* ggml_ext_embedding(GGMLRunnerContext* runner_ctx, + ggml_tensor* weight, + ggml_tensor* input_ids, + ggml_tensor* weight_scale, + int group_size) { + // input_ids: [N, n_token] + + // There are issues with ggml batch inference, so we are expanding it here first. + // TODO: fix ggml batch inference + int64_t n = input_ids->ne[1]; + input_ids = ggml_reshape_1d(runner_ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]); + + ggml_tensor* embedding; + if (weight->type == GGML_TYPE_I8) { + GGML_ASSERT(input_ids->type == GGML_TYPE_I32); + ggml_tensor* args[] = {weight, input_ids, weight_scale}; + embedding = ggml_custom_4d(runner_ctx->ggml_ctx, GGML_TYPE_F32, weight->ne[0], input_ids->ne[0], 1, 1, + args, 3, ggml_ext_get_rows_i8, GGML_N_TASKS_MAX, reinterpret_cast(static_cast(group_size))); + } else { + input_ids = ggml_reshape_3d(runner_ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); + embedding = ggml_get_rows(runner_ctx->ggml_ctx, weight, input_ids); + } + embedding = ggml_reshape_3d(runner_ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); + + // [N, n_token, weight->ne[0]] + return embedding; +} + +ggml_tensor* ggml_ext_prepare_linear_weight(GGMLRunnerContext* runner_ctx, + ggml_tensor* w, + ggml_tensor* x, + bool force_prec_f32) { +#ifndef SD_USE_UPSTREAM_GGML + if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { + bool supports_fp8_matmul = false; + if (runner_ctx->backend != nullptr) { + ggml_tensor* fp8_matmul = ggml_mul_mat(runner_ctx->ggml_ctx, w, x); + if (force_prec_f32) { + ggml_prec_set_acc(fp8_matmul, GGML_PREC_F32); + } + supports_fp8_matmul = ggml_backend_supports_op(runner_ctx->backend, fp8_matmul) && !sd_backend_is_cpu(runner_ctx->backend); + } + if (!supports_fp8_matmul) { + w = ggml_cast(runner_ctx->ggml_ctx, w, GGML_TYPE_BF16); + } + } +#endif + return w; +} + +ggml_tensor* ggml_ext_prepare_i8_input(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + if (x->type != GGML_TYPE_F32 && x->type != GGML_TYPE_I8) { + x = ggml_ext_cast_f32(runner_ctx, x); + } + return ggml_ext_cont(runner_ctx, x); +} + +#ifndef SD_USE_UPSTREAM_GGML +ggml_tensor* ggml_ext_linear_w4_convrot(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* ws, + ggml_tensor* sc, + ggml_tensor* sr, + ggml_tensor* bias, + SDW4Kind kind, + int group, + float scale) { + auto ctx = runner_ctx->ggml_ctx; + x = ggml_ext_prepare_i8_input(runner_ctx, x); + if (scale != 1.f) { + x = ggml_ext_scale(runner_ctx, x, scale); + } + auto out = ggml_ext_mul_mat_w4_convrot(runner_ctx, w, x, ws, sc, sr, scale == 1.f ? bias : nullptr, kind, group); + GGML_ASSERT(out); + if (scale != 1.f) { + out = ggml_ext_scale(runner_ctx, out, 1.f / scale); + if (bias) { + out = ggml_add_inplace(ctx, out, bias); + } + } + return out; +} +#endif + +ggml_tensor* ggml_ext_conv_2d_grouped(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int groups, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1, + bool direct, + bool circular_x, + bool circular_y, + float scale) { + auto ctx = runner_ctx->ggml_ctx; + GGML_ASSERT(groups > 0 && x->ne[2] % groups == 0 && w->ne[3] % groups == 0); + if (groups == 1) { + return ggml_ext_conv_2d(runner_ctx, x, w, b, + s0, s1, p0, p1, d0, d1, + direct, circular_x, circular_y, scale); + } + if (groups == x->ne[2] && groups == w->ne[3]) { + auto out = direct ? ggml_conv_2d_dw_direct(ctx, w, x, s0, s1, p0, p1, d0, d1) + : ggml_conv_2d_dw(ctx, w, x, s0, s1, p0, p1, d0, d1); + if (b) { + out = ggml_add_inplace(ctx, out, ggml_reshape_4d(ctx, b, 1, 1, b->ne[0], 1)); + } + return out; + } + const int64_t ic_g = x->ne[2] / groups, oc_g = w->ne[3] / groups; + std::vector outputs; + for (int i = 0; i < groups; ++i) { + auto xi = ggml_view_4d(ctx, x, x->ne[0], x->ne[1], ic_g, x->ne[3], x->nb[1], x->nb[2], x->nb[3], i * ic_g * x->nb[2]); + auto wi = ggml_view_4d(ctx, w, w->ne[0], w->ne[1], w->ne[2], oc_g, w->nb[1], w->nb[2], w->nb[3], i * oc_g * w->nb[3]); + auto bi = b ? ggml_view_1d(ctx, b, oc_g, i * oc_g * b->nb[0]) : nullptr; + outputs.push_back(ggml_ext_conv_2d(runner_ctx, xi, wi, bi, + s0, s1, p0, p1, d0, d1, + direct, circular_x, circular_y, scale)); + } + return ggml_ext_vec_concat(runner_ctx, outputs, 2); +} diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index 791155780f..f8c14c0ee6 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -10,33 +10,24 @@ struct GGMLRunnerContext; +constexpr uint32_t GGML_PRECISION_F32_VULKAN = 1u << 0; +constexpr uint32_t GGML_PRECISION_F32_ROCM = 1u << 1; + #define EPS 1e-05f static_assert(GGML_MAX_NAME >= 160, "GGML_MAX_NAME must be at least 160"); -// n-mode tensor-matrix product -// example: 2-mode product -// A: [ne03, k, ne01, ne00] -// B: k rows, m columns => [k, m] -// result is [ne03, m, ne01, ne00] -ggml_tensor* ggml_ext_mul_n_mode(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0); - -// Kronecker product -// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10] -ggml_tensor* ggml_ext_kronecker(ggml_context* ctx, ggml_tensor* a, ggml_tensor* b); - -ggml_tensor* ggml_ext_cont(ggml_context* ctx, - ggml_tensor* x); +ggml_tensor* ggml_ext_cont(GGMLRunnerContext* runner_ctx, ggml_tensor* x); // torch like permute -ggml_tensor* ggml_ext_torch_permute(ggml_context* ctx, +ggml_tensor* ggml_ext_torch_permute(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int axis0, int axis1, int axis2, int axis3); -ggml_tensor* ggml_ext_slice(ggml_context* ctx, +ggml_tensor* ggml_ext_slice(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int dim, int64_t start, @@ -44,48 +35,38 @@ ggml_tensor* ggml_ext_slice(ggml_context* ctx, bool cont = true); // example: [N, 3*C, H, W] => ([N, C, H, W], [N, C, H, W], [N, C, H, W]) -std::vector ggml_ext_chunk(ggml_context* ctx, +std::vector ggml_ext_chunk(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int num, int64_t dim, bool cont = true); -ggml_tensor* ggml_ext_silu_act(ggml_context* ctx, ggml_tensor* x, bool gate_first = true); +ggml_tensor* ggml_ext_silu_act(GGMLRunnerContext* runner_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); +ggml_tensor* ggml_ext_situ(GGMLRunnerContext* runner_ctx, ggml_tensor* hidden, ggml_tensor* gate = nullptr); -ggml_tensor* ggml_ext_scale(ggml_context* ctx, - ggml_tensor* x, - float factor, - bool inplace = false); +ggml_tensor* ggml_ext_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* x, 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* ggml_ext_modulate(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* shift, ggml_tensor* scale); -ggml_tensor* ggml_ext_gelu(ggml_context* ctx, - ggml_tensor* x, - bool inplace = false); +ggml_tensor* ggml_ext_gelu(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace = false); -ggml_tensor* ggml_ext_gelu_quick(ggml_context* ctx, - ggml_tensor* x, - bool inplace = false); +ggml_tensor* ggml_ext_gelu_quick(GGMLRunnerContext* runner_ctx, ggml_tensor* x, bool inplace = false); -ggml_tensor* ggml_ext_linear(ggml_context* ctx, +ggml_tensor* ggml_ext_linear(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, bool force_prec_f32 = false, float scale = 1.f); -ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, +ggml_tensor* ggml_ext_linear_i8_tensorwise(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* weight_scale, @@ -93,8 +74,7 @@ ggml_tensor* ggml_ext_linear_i8_tensorwise(ggml_context* ctx, int convrot_group_size, float scale = 1.f); -ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, - ggml_backend_t backend, +ggml_tensor* ggml_ext_pad_ext(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int lp0, int rp0, @@ -107,7 +87,7 @@ ggml_tensor* ggml_ext_pad_ext(ggml_context* ctx, bool circular_x = false, bool circular_y = false); -ggml_tensor* ggml_ext_pad(ggml_context* ctx, +ggml_tensor* ggml_ext_pad(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int p0, int p1, @@ -118,7 +98,7 @@ ggml_tensor* ggml_ext_pad(ggml_context* ctx, // ggml layout: x [L, IC, N], w [K, IC/groups, OC], b [OC], result [OL, OC, N]. // force_prec_f32 keeps both input patches and weights in F32. -ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, +ggml_tensor* ggml_ext_conv_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -132,7 +112,7 @@ ggml_tensor* ggml_ext_conv_1d(ggml_context* ctx, // x: [N, IC, IH, IW] // b: [OC,] // result: [N, OC, OH, OW] -ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, +ggml_tensor* ggml_ext_conv_2d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -151,8 +131,7 @@ ggml_tensor* ggml_ext_conv_2d(ggml_context* ctx, // x: [N, IC, IH, IW] // b: [OC,] // result: [N*OC, OD, OH, OW] -ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, - ggml_backend_t backend, +ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, @@ -168,65 +147,41 @@ ggml_tensor* ggml_ext_conv_3d(ggml_context* ctx, int d2 = 1, bool force_prec_f32 = false, bool direct = false, - float scale = 1.f); - -// w: [OC，IC, KD, 1 * 1] -// x: [N, IC, ID, IH*IW] -// b: [OC,] -// result: [N, OC, OD, OH*OW] -ggml_tensor* ggml_ext_conv_3d_nx1x1(ggml_context* ctx, - ggml_tensor* x, - ggml_tensor* w, - ggml_tensor* b, - int s2 = 1, - int p2 = 1, - int d2 = 1); + float scale = 1.f, + bool allow_conv2d = false); // qkv: [N, L, 3*C] // return: ([N, L, C], [N, L, C], [N, L, C]) -std::vector split_qkv(ggml_context* ctx, - ggml_tensor* qkv); +std::vector ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv); // qkv: [N, 3*C, H, W] // return: ([N, C, H, W], [N, C, H, W], [N, C, H, W]) -std::vector split_image_qkv(ggml_context* ctx, - ggml_tensor* qkv); +std::vector ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv); // Constant and cast helpers require the built-in tensors initialized by GGMLRunner. -ggml_tensor* ggml_ext_full(ggml_context* ctx, +ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx, float value, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3); -ggml_tensor* ggml_ext_zeros(ggml_context* ctx, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3); +ggml_tensor* ggml_ext_zeros(GGMLRunnerContext* runner_ctx, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3); -ggml_tensor* ggml_ext_zeros_like(ggml_context* ctx, - ggml_tensor* x); +ggml_tensor* ggml_ext_zeros_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x); -ggml_tensor* ggml_ext_ones(ggml_context* ctx, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3); +ggml_tensor* ggml_ext_ones(GGMLRunnerContext* runner_ctx, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3); -ggml_tensor* ggml_ext_ones_like(ggml_context* ctx, - ggml_tensor* x); +ggml_tensor* ggml_ext_ones_like(GGMLRunnerContext* runner_ctx, ggml_tensor* x); -ggml_tensor* ggml_ext_cast_f32(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* a); +ggml_tensor* ggml_ext_cast_f32(GGMLRunnerContext* runner_ctx, ggml_tensor* a); // q: [N, L_q, C(n_head*d_head)] or [N*n_head, L_q, d_head] // k: [N, L_k, n_kv_head*d_head] or [N*n_kv_head, L_k, d_head] // v: [N, L_k, n_kv_head*d_head] or [N, L_k, n_kv_head, d_head] // mask: [N, L_q, L_k] // return: [N, L_q, C] -ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, - ggml_backend_t backend, +ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* runner_ctx, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, @@ -235,31 +190,138 @@ ggml_tensor* ggml_ext_attention_ext(ggml_context* ctx, bool skip_reshape = false, bool flash_attn = false, float kv_scale = 1.0f, - bool sage_attn = false, bool* used_flash_attn = nullptr); -ggml_tensor* ggml_ext_layer_norm(ggml_context* ctx, +ggml_tensor* ggml_ext_layer_norm(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, float eps = EPS); -ggml_tensor* ggml_ext_group_norm(ggml_context* ctx, +ggml_tensor* ggml_ext_group_norm(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, int num_groups = 32, float eps = 1e-6f); -ggml_tensor* ggml_ext_timestep_embedding( - ggml_context* ctx, - ggml_tensor* timesteps, - int dim, - int max_period = 10000, - float time_factor = 1.0f); - -ggml_tensor* ggml_ext_vec_concat(ggml_context* ctx, - std::vector& tensors, - int dim); +ggml_tensor* ggml_ext_timestep_embedding(GGMLRunnerContext* runner_ctx, + ggml_tensor* timesteps, + int dim, + int max_period = 10000, + float time_factor = 1.0f); + +ggml_tensor* ggml_ext_vec_concat(GGMLRunnerContext* runner_ctx, std::vector& tensors, int dim); + +#ifndef SD_USE_UPSTREAM_GGML +ggml_tensor* ggml_ext_mul_mat_i8_tensorwise(GGMLRunnerContext* runner_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_ext_quantize_i8_convrot(GGMLRunnerContext* runner_ctx, ggml_tensor* input, int group_size); + +// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook +// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node. +ggml_tensor* ggml_ext_mul_mat_w4_convrot(GGMLRunnerContext* runner_ctx, + ggml_tensor* weight, + ggml_tensor* input, + ggml_tensor* weight_scale, + ggml_tensor* s_channel, + ggml_tensor* s_rel, + ggml_tensor* bias, + SDW4Kind kind, + int group_size); + +// 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_ext_sage_attn(GGMLRunnerContext* runner_ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + float scale, + ggml_sage_attn_mode mode); + +// Noncausal Sol-Attn: contiguous F32 Q/K/V [128, tokens, heads, batch], output [128, heads, tokens, batch]. +ggml_tensor* ggml_ext_sol_attn(GGMLRunnerContext* runner_ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + float scale, + float tau); + +// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes]. +// Returns Q's layout, or nullptr when the backend has no compatible kernel. +ggml_tensor* ggml_ext_window_attn(GGMLRunnerContext* runner_ctx, + ggml_tensor* q, + ggml_tensor* k, + ggml_tensor* v, + const SDWindowAttentionParams& params); +#endif + +// x: F32 [head_dim, heads, tokens, batch], pe: F32 [2, 2, rot_dim/2, tokens]. +// Returns [head_dim, tokens, heads*batch]; channels beyond rot_dim are unchanged. +ggml_tensor* ggml_ext_rope_apply(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* pe, + bool interleaved = true); + +ggml_tensor* ggml_ext_repeat_4d(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + +ggml_tensor* ggml_ext_repeat(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* shape); + +ggml_tensor* ggml_ext_embedding(GGMLRunnerContext* runner_ctx, + ggml_tensor* weight, + ggml_tensor* input_ids, + ggml_tensor* weight_scale = nullptr, + int group_size = 0); + +ggml_tensor* ggml_ext_prepare_linear_weight(GGMLRunnerContext* runner_ctx, + ggml_tensor* w, + ggml_tensor* x, + bool force_prec_f32); + +ggml_tensor* ggml_ext_prepare_i8_input(GGMLRunnerContext* runner_ctx, ggml_tensor* x); + +bool ggml_ext_force_prec_f32(GGMLRunnerContext* runner_ctx, + bool requested, + uint32_t flags); + +#ifndef SD_USE_UPSTREAM_GGML +ggml_tensor* ggml_ext_linear_w4_convrot(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* ws, + ggml_tensor* sc, + ggml_tensor* sr, + ggml_tensor* bias, + SDW4Kind kind, + int group, + float scale = 1.f); +#endif + +ggml_tensor* ggml_ext_conv_2d_grouped(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* w, + ggml_tensor* b, + int groups, + int s0 = 1, + int s1 = 1, + int p0 = 0, + int p1 = 0, + int d0 = 1, + int d1 = 1, + bool direct = false, + bool circular_x = false, + bool circular_y = false, + float scale = 1.f); #endif // __SD_CORE_GGML_EXTEND_H__ diff --git a/src/core/ggml_extensions.cpp b/src/core/ggml_extensions.cpp index dfe5bca402..1e20250634 100644 --- a/src/core/ggml_extensions.cpp +++ b/src/core/ggml_extensions.cpp @@ -131,20 +131,12 @@ bool sd_ggml_rope_describe(const ggml_tensor* node) { 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) { +bool sd_ggml_valid_convrot_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; @@ -155,15 +147,17 @@ static bool packed_height(int64_t k, int64_t rows, int64_t* height) { *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))) + (op.group_size != 0 && (!sd_ggml_valid_convrot_group(op.group_size) || w->ne[0] % op.group_size))) return false; int64_t rows = 1; for (int i = 0; i < 3; ++i) { @@ -177,7 +171,7 @@ static bool i8_shape(const SDI8Operation& op, const int64_t* shape) { 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 || + if (x->type != GGML_TYPE_I8 || !sd_ggml_valid_convrot_group(op.group_size) || op.input_scale || x->ne[1] != height || x->ne[2] != 1 || x->ne[3] != 1) return false; } else { @@ -191,14 +185,15 @@ static bool i8_shape(const SDI8Operation& op, const int64_t* shape) { } 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; + SDI8Attributes 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}; + SDI8Operation op = {inputs[0], inputs[1], nullptr, nullptr, nullptr, attr.group_size, (attr.flags & SD_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)}) { + for (auto entry : {std::make_pair(SD_I8_INPUT_SCALE, &op.input_scale), std::make_pair(SD_I8_WEIGHT_SCALE, &op.weight_scale), std::make_pair(SD_I8_BIAS, &op.bias)}) { if (attr.flags & entry.first) { if (index >= count) return false; @@ -212,10 +207,11 @@ static bool i8_infer(ggml_tensor* const* inputs, int count, const void* attribut 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) + if (count != 1 || inputs[0]->type != GGML_TYPE_F32 || !ggml_is_contiguous(inputs[0]) || !sd_ggml_valid_convrot_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; @@ -227,38 +223,42 @@ static bool convrot_infer(ggml_tensor* const* inputs, int count, const void* att 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; }(); + static const auto op = [] { ggml_ext_op_desc d = {GGML_EXT_ABI_VERSION, "sd.linear_i8.v1", sizeof(SDI8Attributes), 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; + SDI8Attributes 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; + op->packed = (attr.flags & SD_I8_PACKED) != 0; int index = 2; - if (attr.flags & I8_INPUT_SCALE) + if (attr.flags & SD_I8_INPUT_SCALE) op->input_scale = node->src[index++]; - if (attr.flags & I8_WEIGHT_SCALE) + if (attr.flags & SD_I8_WEIGHT_SCALE) op->weight_scale = node->src[index++]; - if (attr.flags & I8_BIAS) + if (attr.flags & SD_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)) + if (((attr.flags & SD_I8_INPUT_SCALE) && !op->input_scale) || ((attr.flags & SD_I8_WEIGHT_SCALE) && !op->weight_scale) || ((attr.flags & SD_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; @@ -269,48 +269,11 @@ bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) { 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); -} - -struct W4Attributes { - int64_t shape[3]; - SDW4Kind kind; - int32_t has_bias; -}; static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { if (!attributes || count < 3) return false; - W4Attributes attr; + SDW4Attributes attr; std::memcpy(&attr, attributes, sizeof(attr)); const bool codebook = attr.kind == SDW4Kind::W4A8; if ((!codebook && attr.kind != SDW4Kind::W4A4) || (attr.has_bias != 0 && attr.has_bias != 1) || @@ -348,7 +311,7 @@ static bool w4_infer(ggml_tensor* const* inputs, int count, const void* attribut ggml_ext_op_t sd_ggml_w4_op() { static const auto op = [] { - const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(W4Attributes), w4_infer}; + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.linear_w4_convrot.v1", sizeof(SDW4Attributes), w4_infer}; auto result = ggml_ext_register_op(&desc); GGML_ASSERT(result); return result; @@ -359,7 +322,7 @@ ggml_ext_op_t sd_ggml_w4_op() { bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) { if (!node || ggml_ext_op_get(node) != sd_ggml_w4_op() || !ggml_is_contiguous(node)) return false; - W4Attributes attr; + SDW4Attributes attr; std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(attr)); const bool codebook = attr.kind == SDW4Kind::W4A8; const int count = (codebook ? 5 : 3) + (attr.has_bias != 0); @@ -372,30 +335,6 @@ bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* op) { return true; } -ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* sc, ggml_tensor* sr, ggml_tensor* bias, SDW4Kind kind, int group) { - if (!w || !x || group < 0 || (group && !valid_group(group))) - return nullptr; - const int quant_group = group ? group : 1; - if (x->type == GGML_TYPE_F32) - x = ggml_quantize_i8_convrot(ctx, x, quant_group); - int packed_group; - if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != quant_group) - return nullptr; - const auto* logical = x->src[0]; - W4Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, kind, bias ? 1 : 0}; - ggml_tensor* inputs[6] = {w, x, ws}; - int count = 3; - if (kind == SDW4Kind::W4A8) { - inputs[count++] = sc; - inputs[count++] = sr; - } else if (sc || sr) { - return nullptr; - } - if (bias) - inputs[count++] = bias; - return ggml_ext_op_create(ctx, sd_ggml_w4_op(), inputs, count, &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; @@ -443,12 +382,6 @@ bool sd_ggml_sage_attn_describe(const ggml_tensor* node, SDSageAttentionParams* 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; @@ -495,13 +428,10 @@ bool sd_ggml_sol_attn_describe(const ggml_tensor* node, SDSolAttentionParams* pa 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); -} - -static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void* attributes, ggml_ext_tensor_desc* out) { +bool sd_ggml_window_attn_infer(ggml_tensor* const* inputs, + int count, + const void* attributes, + ggml_ext_tensor_desc* out) { if (count != 3 || !attributes) return false; SDWindowAttentionParams p; @@ -531,7 +461,7 @@ static bool window_attn_infer(ggml_tensor* const* inputs, int count, const void* ggml_ext_op_t sd_ggml_window_attn_op() { static const auto op = [] { - const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), window_attn_infer}; + const ggml_ext_op_desc desc = {GGML_EXT_ABI_VERSION, "sd.window_attn.v1", sizeof(SDWindowAttentionParams), sd_ggml_window_attn_infer}; auto result = ggml_ext_register_op(&desc); GGML_ASSERT(result); return result; @@ -544,22 +474,12 @@ bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionPara return false; const auto* attributes = ggml_ext_op_attributes(node); ggml_ext_tensor_desc out; - if (!window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne))) + if (!sd_ggml_window_attn_infer(node->src, 3, attributes, &out) || node->type != out.type || std::memcmp(node->ne, out.ne, sizeof(out.ne))) return false; std::memcpy(params, attributes, sizeof(*params)); return true; } -ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params) { - ggml_tensor* inputs[] = {q, k, v}; - ggml_ext_tensor_desc out; - if (!backend || !window_attn_infer(inputs, 3, ¶ms, &out)) - return nullptr; - sd_ggml_register_extensions(backend); - auto result = ggml_ext_op_create(ctx, sd_ggml_window_attn_op(), inputs, 3, ¶ms); - return ggml_backend_supports_op(backend, result) ? result : nullptr; -} - void sd_ggml_register_extensions(ggml_backend_t backend) { if (!backend) { return; @@ -583,13 +503,6 @@ void sd_ggml_register_extensions(ggml_backend_t backend) { 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 index 87a51a7e5e..3145c6fbb4 100644 --- a/src/core/ggml_extensions.h +++ b/src/core/ggml_extensions.h @@ -9,18 +9,24 @@ void sd_ggml_register_extensions(ggml_backend_t backend); #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 SDI8Attributes { + int64_t shape[3]; + int32_t group_size; + int32_t flags; +}; +enum SDI8Flags { SD_I8_PACKED = 1, + SD_I8_INPUT_SCALE = 2, + SD_I8_WEIGHT_SCALE = 4, + SD_I8_BIAS = 8 }; struct SDI8Operation { const ggml_tensor* weight; @@ -32,16 +38,25 @@ struct SDI8Operation { bool packed; }; +bool sd_ggml_valid_convrot_group(int group); + 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 class SDW4Kind : int32_t { W4A4 = 1, W4A8 = 2 }; +struct SDW4Attributes { + int64_t shape[3]; + SDW4Kind kind; + int32_t has_bias; +}; + struct SDW4Operation { const ggml_tensor* weight; const ggml_tensor* input; @@ -53,10 +68,8 @@ struct SDW4Operation { }; ggml_ext_op_t sd_ggml_w4_op(); + bool sd_ggml_w4_describe(const ggml_tensor* node, SDW4Operation* operation); -// Packed I8 [K/2, N], low nibble first. W4A8 uses a 16-entry F32 codebook -// and F8_E4M3 [K/16, N] relative scales. F32 inputs get an explicit quantize node. -ggml_tensor* ggml_mul_mat_w4_convrot(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* s_channel, ggml_tensor* s_rel, ggml_tensor* bias, SDW4Kind kind, int group_size); enum ggml_sage_attn_mode { GGML_SAGE_ATTN_AUTO = 0, @@ -70,10 +83,8 @@ struct SDSageAttentionParams { }; 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; @@ -86,16 +97,19 @@ struct SDWindowAttentionParams { int32_t preserve_boundary; }; +bool sd_ggml_window_attn_infer(ggml_tensor* const* inputs, + int count, + const void* attributes, + ggml_ext_tensor_desc* out); + ggml_ext_op_t sd_ggml_window_attn_op(); + bool sd_ggml_window_attn_describe(const ggml_tensor* node, SDWindowAttentionParams* params); -// Contiguous F32 Q [D*heads, W*H], K/V [D*kv_heads, W*H*planes]. -// Returns Q's layout, or nullptr when the backend has no compatible kernel. -ggml_tensor* sd_ggml_window_attn(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* q, ggml_tensor* k, ggml_tensor* v, const SDWindowAttentionParams& params); 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_runner.cpp b/src/core/ggml_runner.cpp index 770f46b219..39d5429744 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -13,22 +13,6 @@ using namespace sd; -ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, - ggml_tensor* q, - ggml_tensor* k, - ggml_tensor* v, - int64_t n_head, - ggml_tensor* mask, - bool skip_reshape, - bool flash_attn, - float kv_scale, - bool* used_flash_attn) { - if (ctx->attn_scale > 0.f) { - kv_scale = ctx->attn_scale; - } - return ggml_ext_attention_ext(ctx->ggml_ctx, ctx->backend, q, k, v, n_head, mask, skip_reshape, flash_attn, kv_scale, ctx->sage_attn_enabled, used_flash_attn); -} - void GGMLRunner::alloc_params_ctx() { ggml_init_params params; params.mem_size = static_cast(MAX_PARAMS_TENSOR_NUM * ggml_tensor_overhead()); diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index 5c21b1daeb..d300be7681 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -23,6 +23,8 @@ #define MAX_PARAMS_TENSOR_NUM 32768 #define MAX_GRAPH_SIZE 327680 +struct GGMLRunnerContext; + struct WeightAdapter { struct ForwardParams { enum class op_type_t { @@ -46,22 +48,26 @@ struct WeightAdapter { float scale = 1.f; } conv2d; }; - virtual ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) = 0; - virtual ggml_tensor* forward_with_lora(ggml_context* ctx, - ggml_backend_t backend, + + virtual ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx, + ggml_tensor* weight, + const std::string& weight_name) = 0; + + virtual ggml_tensor* forward_with_lora(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, const std::string& prefix, - ForwardParams forward_params) = 0; - virtual ggml_tensor* add_lora_to_output(ggml_context* ctx, - ggml_backend_t backend, + ForwardParams forward_params) = 0; + + virtual ggml_tensor* add_lora_to_output(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* output, const std::string& prefix, - ForwardParams forward_params) = 0; - virtual size_t get_extra_graph_size() = 0; + ForwardParams forward_params) = 0; + + virtual size_t get_extra_graph_size() = 0; }; struct GGMLRunnerContext { @@ -124,17 +130,6 @@ struct GGMLRunnerContext { } }; -ggml_tensor* ggml_ext_attention_ext(GGMLRunnerContext* ctx, - ggml_tensor* q, - ggml_tensor* k, - ggml_tensor* v, - int64_t n_head, - ggml_tensor* mask = nullptr, - bool skip_reshape = false, - bool flash_attn = false, - float kv_scale = 1.f, - bool* used_flash_attn = nullptr); - struct GGMLRunner { private: std::map logged_compute_bytes_; @@ -142,11 +137,15 @@ struct GGMLRunner { ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; 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, 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: diff --git a/src/core/probe/flash_attn.cpp b/src/core/probe/flash_attn.cpp index 3dcb95e826..984de1c676 100644 --- a/src/core/probe/flash_attn.cpp +++ b/src/core/probe/flash_attn.cpp @@ -9,6 +9,7 @@ #include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" #include "core/probe/backend_probe.h" #include "core/util.h" @@ -41,8 +42,11 @@ static void build_flash_attention_probe(BackendProbeGraph& graph, ggml_backend_t auto k = ggml_dup_tensor(graph.ctx, q); auto v = ggml_dup_tensor(graph.ctx, q); graph.inputs = {q, k, v}; + GGMLRunnerContext runner_ctx; + runner_ctx.ggml_ctx = graph.ctx; + runner_ctx.backend = backend; bool used_flash = false; - auto output = ggml_ext_attention_ext(graph.ctx, backend, q, k, v, heads, nullptr, false, flash, 1.f, false, &used_flash); + auto output = ggml_ext_attention_ext(&runner_ctx, q, k, v, heads, nullptr, false, flash, 1.f, &used_flash); graph.output = flash && !used_flash ? nullptr : output; } diff --git a/src/model/adapter/lora.hpp b/src/model/adapter/lora.hpp index 7b9e8ab82d..7d71b34f60 100644 --- a/src/model/adapter/lora.hpp +++ b/src/model/adapter/lora.hpp @@ -218,7 +218,8 @@ struct LoraModel : public GGMLRunner { lora_tensors = std::move(new_lora_tensors); } - ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) { + ggml_tensor* get_lora_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) { + auto ctx = runner_ctx->ggml_ctx; ggml_tensor* updown = nullptr; int index = 0; while (true) { @@ -241,21 +242,21 @@ struct LoraModel : public GGMLRunner { auto iter = lora_tensors.find(lora_up_name); if (iter != lora_tensors.end()) { - lora_up = ggml_ext_cast_f32(ctx, backend, iter->second); + lora_up = ggml_ext_cast_f32(runner_ctx, iter->second); if (swapped_lora_up_tensors.count(lora_up_name)) { - auto halves = ggml_ext_chunk(ctx, lora_up, 2, 1); + auto halves = ggml_ext_chunk(runner_ctx, lora_up, 2, 1); lora_up = ggml_concat(ctx, halves[1], halves[0], 1); } } iter = lora_tensors.find(lora_mid_name); if (iter != lora_tensors.end()) { - lora_mid = ggml_ext_cast_f32(ctx, backend, iter->second); + lora_mid = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(lora_down_name); if (iter != lora_tensors.end()) { - lora_down = ggml_ext_cast_f32(ctx, backend, iter->second); + lora_down = ggml_ext_cast_f32(runner_ctx, iter->second); } if (lora_up == nullptr || lora_down == nullptr) { @@ -288,8 +289,8 @@ struct LoraModel : public GGMLRunner { } scale_value *= multiplier; - auto curr_updown = ggml_ext_merge_lora(ctx, lora_down, lora_up, lora_mid); - curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true); + auto curr_updown = ggml_ext_merge_lora(runner_ctx, lora_down, lora_up, lora_mid); + curr_updown = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true); if (updown == nullptr) { updown = curr_updown; @@ -302,7 +303,8 @@ struct LoraModel : public GGMLRunner { return updown; } - ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) { + ggml_tensor* get_raw_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) { + auto ctx = runner_ctx->ggml_ctx; ggml_tensor* updown = nullptr; int index = 0; while (true) { @@ -319,7 +321,7 @@ struct LoraModel : public GGMLRunner { auto iter = lora_tensors.find(diff_name); if (iter != lora_tensors.end()) { - curr_updown = ggml_ext_cast_f32(ctx, backend, iter->second); + curr_updown = ggml_ext_cast_f32(runner_ctx, iter->second); } else { break; } @@ -329,7 +331,7 @@ struct LoraModel : public GGMLRunner { float scale_value = 1.0f; scale_value *= multiplier; - curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true); + curr_updown = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true); if (updown == nullptr) { updown = curr_updown; @@ -342,7 +344,8 @@ struct LoraModel : public GGMLRunner { return updown; } - ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) { + ggml_tensor* get_loha_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) { + auto ctx = runner_ctx->ggml_ctx; ggml_tensor* updown = nullptr; int index = 0; while (true) { @@ -370,17 +373,17 @@ struct LoraModel : public GGMLRunner { auto iter = lora_tensors.find(hada_1_down_name); if (iter != lora_tensors.end()) { - hada_1_down = ggml_ext_cast_f32(ctx, backend, iter->second); + hada_1_down = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(hada_1_up_name); if (iter != lora_tensors.end()) { - hada_1_up = ggml_ext_cast_f32(ctx, backend, iter->second); + hada_1_up = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(hada_1_mid_name); if (iter != lora_tensors.end()) { - hada_1_mid = ggml_ext_cast_f32(ctx, backend, iter->second); + hada_1_mid = ggml_ext_cast_f32(runner_ctx, iter->second); if (hada_1_up != nullptr) { hada_1_up = ggml_cont(ctx, ggml_transpose(ctx, hada_1_up)); } @@ -388,17 +391,17 @@ struct LoraModel : public GGMLRunner { iter = lora_tensors.find(hada_2_down_name); if (iter != lora_tensors.end()) { - hada_2_down = ggml_ext_cast_f32(ctx, backend, iter->second); + hada_2_down = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(hada_2_up_name); if (iter != lora_tensors.end()) { - hada_2_up = ggml_ext_cast_f32(ctx, backend, iter->second); + hada_2_up = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(hada_2_mid_name); if (iter != lora_tensors.end()) { - hada_2_mid = ggml_ext_cast_f32(ctx, backend, iter->second); + hada_2_mid = ggml_ext_cast_f32(runner_ctx, iter->second); if (hada_2_up != nullptr) { hada_2_up = ggml_cont(ctx, ggml_transpose(ctx, hada_2_up)); } @@ -435,10 +438,10 @@ struct LoraModel : public GGMLRunner { } scale_value *= multiplier; - ggml_tensor* updown_1 = ggml_ext_merge_lora(ctx, hada_1_down, hada_1_up, hada_1_mid); - ggml_tensor* updown_2 = ggml_ext_merge_lora(ctx, hada_2_down, hada_2_up, hada_2_mid); + ggml_tensor* updown_1 = ggml_ext_merge_lora(runner_ctx, hada_1_down, hada_1_up, hada_1_mid); + ggml_tensor* updown_2 = ggml_ext_merge_lora(runner_ctx, hada_2_down, hada_2_up, hada_2_mid); auto curr_updown = ggml_mul_inplace(ctx, updown_1, updown_2); - curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true); + curr_updown = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true); if (updown == nullptr) { updown = curr_updown; } else { @@ -449,7 +452,8 @@ struct LoraModel : public GGMLRunner { return updown; } - ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, ggml_context* ctx, ggml_backend_t backend) { + ggml_tensor* get_lokr_weight_diff(const std::string& model_tensor_name, GGMLRunnerContext* runner_ctx) { + auto ctx = runner_ctx->ggml_ctx; ggml_tensor* updown = nullptr; int index = 0; while (true) { @@ -476,24 +480,24 @@ struct LoraModel : public GGMLRunner { auto iter = lora_tensors.find(lokr_w1_name); if (iter != lora_tensors.end()) { - lokr_w1 = ggml_ext_cast_f32(ctx, backend, iter->second); + lokr_w1 = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(lokr_w2_name); if (iter != lora_tensors.end()) { - lokr_w2 = ggml_ext_cast_f32(ctx, backend, iter->second); + lokr_w2 = ggml_ext_cast_f32(runner_ctx, iter->second); } int64_t rank = 1; if (lokr_w1 == nullptr) { iter = lora_tensors.find(lokr_w1_a_name); if (iter != lora_tensors.end()) { - lokr_w1_a = ggml_ext_cast_f32(ctx, backend, iter->second); + lokr_w1_a = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(lokr_w1_b_name); if (iter != lora_tensors.end()) { - lokr_w1_b = ggml_ext_cast_f32(ctx, backend, iter->second); + lokr_w1_b = ggml_ext_cast_f32(runner_ctx, iter->second); } if (lokr_w1_a == nullptr || lokr_w1_b == nullptr) { @@ -502,18 +506,18 @@ struct LoraModel : public GGMLRunner { rank = lokr_w1_b->ne[ggml_n_dims(lokr_w1_b) - 1]; - lokr_w1 = ggml_ext_merge_lora(ctx, lokr_w1_b, lokr_w1_a); + lokr_w1 = ggml_ext_merge_lora(runner_ctx, lokr_w1_b, lokr_w1_a); } if (lokr_w2 == nullptr) { iter = lora_tensors.find(lokr_w2_a_name); if (iter != lora_tensors.end()) { - lokr_w2_a = ggml_ext_cast_f32(ctx, backend, iter->second); + lokr_w2_a = ggml_ext_cast_f32(runner_ctx, iter->second); } iter = lora_tensors.find(lokr_w2_b_name); if (iter != lora_tensors.end()) { - lokr_w2_b = ggml_ext_cast_f32(ctx, backend, iter->second); + lokr_w2_b = ggml_ext_cast_f32(runner_ctx, iter->second); } if (lokr_w2_a == nullptr || lokr_w2_b == nullptr) { @@ -522,7 +526,7 @@ struct LoraModel : public GGMLRunner { rank = lokr_w2_b->ne[ggml_n_dims(lokr_w2_b) - 1]; - lokr_w2 = ggml_ext_merge_lora(ctx, lokr_w2_b, lokr_w2_a); + lokr_w2 = ggml_ext_merge_lora(runner_ctx, lokr_w2_b, lokr_w2_a); } if (!lokr_w1_a) { @@ -553,8 +557,8 @@ struct LoraModel : public GGMLRunner { scale_value *= multiplier; - auto curr_updown = ggml_ext_kronecker(ctx, lokr_w1, lokr_w2); - curr_updown = ggml_ext_scale(ctx, curr_updown, scale_value, true); + auto curr_updown = ggml_ext_kronecker(runner_ctx, lokr_w1, lokr_w2); + curr_updown = ggml_ext_scale(runner_ctx, curr_updown, scale_value, true); if (updown == nullptr) { updown = curr_updown; @@ -566,29 +570,33 @@ struct LoraModel : public GGMLRunner { return updown; } - ggml_tensor* get_weight_diff(const std::string& model_tensor_name, ggml_backend_t backend, ggml_context* ctx, ggml_tensor* model_tensor, bool with_lora_and_lokr = true) { + ggml_tensor* get_weight_diff(const std::string& model_tensor_name, + GGMLRunnerContext* runner_ctx, + ggml_tensor* model_tensor, + bool with_lora_and_lokr = true) { + auto ctx = runner_ctx->ggml_ctx; // lora ggml_tensor* diff = nullptr; if (with_lora_and_lokr) { - diff = get_lora_weight_diff(model_tensor_name, ctx, backend); + diff = get_lora_weight_diff(model_tensor_name, runner_ctx); } // diff if (diff == nullptr) { - diff = get_raw_weight_diff(model_tensor_name, ctx, backend); + diff = get_raw_weight_diff(model_tensor_name, runner_ctx); } // loha if (diff == nullptr) { - diff = get_loha_weight_diff(model_tensor_name, ctx, backend); + diff = get_loha_weight_diff(model_tensor_name, runner_ctx); } // lokr if (diff == nullptr && with_lora_and_lokr) { - diff = get_lokr_weight_diff(model_tensor_name, ctx, backend); + diff = get_lokr_weight_diff(model_tensor_name, runner_ctx); } if (diff != nullptr) { if (ggml_nelements(diff) < ggml_nelements(model_tensor)) { if (ggml_n_dims(diff) == 2 && ggml_n_dims(model_tensor) == 2 && diff->ne[0] == model_tensor->ne[0]) { LOG_WARN("pad for %s", model_tensor_name.c_str()); - auto pad_tensor = ggml_ext_zeros(ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1); + auto pad_tensor = ggml_ext_zeros(runner_ctx, diff->ne[0], model_tensor->ne[1] - diff->ne[1], 1, 1); diff = ggml_concat(ctx, diff, pad_tensor, 1); } } @@ -619,12 +627,12 @@ struct LoraModel : public GGMLRunner { return diff; } - ggml_tensor* get_out_diff(ggml_context* ctx, - ggml_backend_t backend, + ggml_tensor* get_out_diff(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* model_weight, WeightAdapter::ForwardParams forward_params, const std::string& model_tensor_name) { + auto ctx = runner_ctx->ggml_ctx; ggml_tensor* out_diff = nullptr; int index = 0; @@ -712,7 +720,7 @@ struct LoraModel : public GGMLRunner { } scale_value *= multiplier; - auto curr_out_diff = ggml_ext_lokr_forward(ctx, backend, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value); + auto curr_out_diff = ggml_ext_lokr_forward(runner_ctx, x, lokr_w1, lokr_w1_a, lokr_w1_b, lokr_w2, lokr_w2_a, lokr_w2_b, is_conv2d, forward_params.conv2d, scale_value); if (out_diff == nullptr) { out_diff = curr_out_diff; } else { @@ -761,7 +769,7 @@ struct LoraModel : public GGMLRunner { if (iter != lora_tensors.end()) { lora_up = iter->second; if (swapped_lora_up_tensors.count(lora_up_name)) { - auto halves = ggml_ext_chunk(ctx, lora_up, 2, 1); + auto halves = ggml_ext_chunk(runner_ctx, lora_up, 2, 1); lora_up = ggml_concat(ctx, halves[1], halves[0], 1); } if (is_conv2d && lora_up->type != GGML_TYPE_F16) { @@ -850,13 +858,13 @@ struct LoraModel : public GGMLRunner { ggml_tensor* lx; if (!is_conv2d) { - lx = ggml_ext_linear(ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale); + lx = ggml_ext_linear(runner_ctx, x, lora_down, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale); if (lora_mid) { - lx = ggml_ext_linear(ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale); + lx = ggml_ext_linear(runner_ctx, lx, lora_mid, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale); } - lx = ggml_ext_linear(ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale); + lx = ggml_ext_linear(runner_ctx, lx, lora_up, nullptr, forward_params.linear.force_prec_f32, forward_params.linear.scale); } else { // OP_CONV2D - lx = ggml_ext_conv_2d(ctx, + lx = ggml_ext_conv_2d(runner_ctx, x, lora_down, nullptr, @@ -871,7 +879,7 @@ struct LoraModel : public GGMLRunner { forward_params.conv2d.circular_y, forward_params.conv2d.scale); if (lora_mid) { - lx = ggml_ext_conv_2d(ctx, + lx = ggml_ext_conv_2d(runner_ctx, lx, lora_mid, nullptr, @@ -886,7 +894,7 @@ struct LoraModel : public GGMLRunner { forward_params.conv2d.circular_y, forward_params.conv2d.scale); } - lx = ggml_ext_conv_2d(ctx, + lx = ggml_ext_conv_2d(runner_ctx, lx, lora_up, nullptr, @@ -902,7 +910,7 @@ struct LoraModel : public GGMLRunner { forward_params.conv2d.scale); } - auto curr_out_diff = ggml_ext_scale(ctx, lx, scale_value, true); + auto curr_out_diff = ggml_ext_scale(runner_ctx, lx, scale_value, true); if (out_diff == nullptr) { out_diff = curr_out_diff; } else { @@ -951,6 +959,7 @@ struct LoraModel : public GGMLRunner { size_t lora_graph_size = LORA_GRAPH_BASE_SIZE + lora_tensors.size() * 10; ggml_cgraph* gf = ggml_new_graph_custom(compute_ctx, lora_graph_size, false); + auto runner_ctx = get_context(); preprocess_lora_tensors(model_tensor_names); original_tensor_to_final_tensor.clear(); @@ -961,7 +970,7 @@ struct LoraModel : public GGMLRunner { ggml_tensor* model_tensor = it.second; // lora - ggml_tensor* diff = get_weight_diff(model_tensor_name, runtime_backend, compute_ctx, model_tensor); + ggml_tensor* diff = get_weight_diff(model_tensor_name, &runner_ctx, model_tensor); if (diff == nullptr) { continue; } @@ -974,7 +983,7 @@ struct LoraModel : public GGMLRunner { ggml_tensor* final_tensor; if (model_tensor->type != GGML_TYPE_F32 && model_tensor->type != GGML_TYPE_F16) { - final_tensor = ggml_ext_cast_f32(compute_ctx, runtime_backend, model_tensor); + final_tensor = ggml_ext_cast_f32(&runner_ctx, model_tensor); final_tensor = ggml_add_inplace(compute_ctx, final_tensor, diff); final_tensor = ggml_cpy(compute_ctx, final_tensor, model_tensor); } else { @@ -1061,41 +1070,47 @@ struct MultiLoraAdapter : public WeightAdapter { : lora_models(lora_models) { } - ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name, bool with_lora_and_lokr) { + ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx, + ggml_tensor* weight, + const std::string& weight_name, + bool with_lora_and_lokr) { + auto ctx = runner_ctx->ggml_ctx; for (auto& lora_model : lora_models) { - ggml_tensor* diff = lora_model->get_weight_diff(weight_name, backend, ctx, weight, with_lora_and_lokr); + ggml_tensor* diff = lora_model->get_weight_diff(weight_name, runner_ctx, weight, with_lora_and_lokr); if (diff == nullptr) { continue; } if (weight->type != GGML_TYPE_F32 && weight->type != GGML_TYPE_F16) { - weight = ggml_ext_cast_f32(ctx, backend, weight); + weight = ggml_ext_cast_f32(runner_ctx, weight); } weight = ggml_add(ctx, weight, diff); } return weight; } - ggml_tensor* patch_weight(ggml_context* ctx, ggml_backend_t backend, ggml_tensor* weight, const std::string& weight_name) override { - return patch_weight(ctx, backend, weight, weight_name, true); + ggml_tensor* patch_weight(GGMLRunnerContext* runner_ctx, + ggml_tensor* weight, + const std::string& weight_name) override { + return patch_weight(runner_ctx, weight, weight_name, true); } - ggml_tensor* forward_with_lora(ggml_context* ctx, - ggml_backend_t backend, + ggml_tensor* forward_with_lora(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* b, const std::string& prefix, WeightAdapter::ForwardParams forward_params) override { - w = patch_weight(ctx, backend, w, prefix + "weight", false); + auto ctx = runner_ctx->ggml_ctx; + w = patch_weight(runner_ctx, w, prefix + "weight", false); if (b) { - b = patch_weight(ctx, backend, b, prefix + "bias", false); + b = patch_weight(runner_ctx, b, prefix + "bias", false); } ggml_tensor* out; if (forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR) { - out = ggml_ext_linear(ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale); + out = ggml_ext_linear(runner_ctx, x, w, b, forward_params.linear.force_prec_f32, forward_params.linear.scale); } else { // OP_CONV2D - out = ggml_ext_conv_2d(ctx, + out = ggml_ext_conv_2d(runner_ctx, x, w, b, @@ -1111,7 +1126,7 @@ struct MultiLoraAdapter : public WeightAdapter { forward_params.conv2d.scale); } for (auto& lora_model : lora_models) { - ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight"); + ggml_tensor* out_diff = lora_model->get_out_diff(runner_ctx, x, w, forward_params, prefix + "weight"); if (out_diff == nullptr) { continue; } @@ -1120,13 +1135,13 @@ struct MultiLoraAdapter : public WeightAdapter { return out; } - ggml_tensor* add_lora_to_output(ggml_context* ctx, - ggml_backend_t backend, + ggml_tensor* add_lora_to_output(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* w, ggml_tensor* output, const std::string& prefix, WeightAdapter::ForwardParams forward_params) override { + auto ctx = runner_ctx->ggml_ctx; // Packed INT4 stores half the logical input width. Adapter shape checks // and reshapes need the logical weight dimensions, without decoding it. ggml_tensor logical_weight = *w; @@ -1138,10 +1153,10 @@ struct MultiLoraAdapter : public WeightAdapter { w = &logical_weight; } for (auto& lora_model : lora_models) { - ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", backend, ctx, w, false); + ggml_tensor* weight_diff = lora_model->get_weight_diff(prefix + "weight", runner_ctx, w, false); if (weight_diff != nullptr) { GGML_ASSERT(forward_params.op_type == ForwardParams::op_type_t::OP_LINEAR); - ggml_tensor* out_diff = ggml_ext_linear(ctx, + ggml_tensor* out_diff = ggml_ext_linear(runner_ctx, x, weight_diff, nullptr, @@ -1150,7 +1165,7 @@ struct MultiLoraAdapter : public WeightAdapter { output = ggml_add_inplace(ctx, output, out_diff); } - ggml_tensor* out_diff = lora_model->get_out_diff(ctx, backend, x, w, forward_params, prefix + "weight"); + ggml_tensor* out_diff = lora_model->get_out_diff(runner_ctx, x, w, forward_params, prefix + "weight"); if (out_diff != nullptr) { output = ggml_add_inplace(ctx, output, out_diff); } diff --git a/src/model/adapter/lora_ops.cpp b/src/model/adapter/lora_ops.cpp index 5a6f905de9..9629c339db 100644 --- a/src/model/adapter/lora_ops.cpp +++ b/src/model/adapter/lora_ops.cpp @@ -4,11 +4,46 @@ #include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" +#include "core/ggml_runner.h" + +ggml_tensor* ggml_ext_mul_n_mode(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b, int mode) { + auto ctx = runner_ctx->ggml_ctx; + // reshape A + // swap 0th and nth axis + a = ggml_cont(ctx, ggml_permute(ctx, a, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0)); + int64_t ne1 = a->ne[1]; + int64_t ne2 = a->ne[2]; + int64_t ne3 = a->ne[3]; + // make 2D + a = ggml_cont(ctx, ggml_reshape_2d(ctx, a, a->ne[0], (ne3 * ne2 * ne1))); + + ggml_tensor* result = ggml_cont(ctx, ggml_transpose(ctx, ggml_mul_mat(ctx, a, b))); + + // reshape output (same shape as a after permutation except first dim) + result = ggml_reshape_4d(ctx, result, result->ne[0], ne1, ne2, ne3); + // swap back 0th and nth axis + result = ggml_permute(ctx, result, mode, mode != 1 ? 1 : 0, mode != 2 ? 2 : 0, mode != 3 ? 3 : 0); + return result; +} + +ggml_tensor* ggml_ext_kronecker(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b) { + auto ctx = runner_ctx->ggml_ctx; + return ggml_mul(ctx, + ggml_interpolate(ctx, + a, + a->ne[0] * b->ne[0], + a->ne[1] * b->ne[1], + a->ne[2] * b->ne[2], + a->ne[3] * b->ne[3], + GGML_SCALE_MODE_NEAREST), + b); +} -ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, +ggml_tensor* ggml_ext_merge_lora(GGMLRunnerContext* runner_ctx, ggml_tensor* lora_down, ggml_tensor* lora_up, ggml_tensor* lora_mid) { + auto ctx = runner_ctx->ggml_ctx; ggml_tensor* updown; // flat lora tensors to multiply it // ggml_n_dims drops the output dimension when a linear layer has one output. @@ -32,25 +67,25 @@ ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, // lora_down has shape (Rank, In, 1, 1) // lora_up has shape (Rank, Out, 1, 1) // conv layer shape is (3, 3, Out, In) - updown = ggml_ext_mul_n_mode(ctx, ggml_ext_mul_n_mode(ctx, lora_mid, lora_down, 3), lora_up, 2); + updown = ggml_ext_mul_n_mode(runner_ctx, ggml_ext_mul_n_mode(runner_ctx, lora_mid, lora_down, 3), lora_up, 2); updown = ggml_cont(ctx, updown); } return updown; } -ggml_tensor* ggml_ext_lokr_forward( - ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] - ggml_tensor* w1, // Outer C (Full rank) - ggml_tensor* w1a, // Outer A (Low rank part 1) - ggml_tensor* w1b, // Outer B (Low rank part 2) - ggml_tensor* w2, // Inner BA (Full rank) - ggml_tensor* w2a, // Inner A (Low rank part 1) - ggml_tensor* w2b, // Inner B (Low rank part 2) - bool is_conv, - WeightAdapter::ForwardParams::conv2d_params_t conv_params, - float scale) { +ggml_tensor* ggml_ext_lokr_forward(GGMLRunnerContext* runner_ctx, + ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] + ggml_tensor* w1, // Outer C (Full rank) + ggml_tensor* w1a, // Outer A (Low rank part 1) + ggml_tensor* w1b, // Outer B (Low rank part 2) + ggml_tensor* w2, // Inner BA (Full rank) + ggml_tensor* w2a, // Inner A (Low rank part 1) + ggml_tensor* w2b, // Inner B (Low rank part 2) + bool is_conv, + WeightAdapter::ForwardParams::conv2d_params_t conv_params, + float scale) { + auto ctx = runner_ctx->ggml_ctx; + auto backend = runner_ctx->backend; GGML_ASSERT((w1 != nullptr || (w1a != nullptr && w1b != nullptr))); GGML_ASSERT((w2 != nullptr || (w2a != nullptr && w2b != nullptr))); @@ -73,7 +108,6 @@ ggml_tensor* ggml_ext_lokr_forward( if (sd_backend_is(backend, "Vulkan")) { if (batch > 1) { - // no access to backend here, worst case is slightly worse perfs for other backends when built alongside Vulkan backend int max_batch = 65535; int max_batch_uq = max_batch / uq; merge_batch_uq = 1; @@ -121,14 +155,14 @@ ggml_tensor* ggml_ext_lokr_forward( ggml_tensor* hc = ggml_transpose(ctx, hc_t); ggml_tensor* out = ggml_reshape_2d(ctx, ggml_cont(ctx, hc), up * vp, batch); - return ggml_ext_scale(ctx, out, scale); + return ggml_ext_scale(runner_ctx, out, scale); } else { int batch = (int)h->ne[3]; // 1. Reshape input: [W, H, vq*uq, batch] -> [W, H, vq, uq * batch] ggml_tensor* h_split = ggml_reshape_4d(ctx, h, h->ne[0], h->ne[1], vq, uq * batch); if (w2 != nullptr) { - hb = ggml_ext_conv_2d(ctx, h_split, w2, nullptr, + hb = ggml_ext_conv_2d(runner_ctx, h_split, w2, nullptr, conv_params.s0, conv_params.s1, conv_params.p0, @@ -154,7 +188,7 @@ ggml_tensor* ggml_ext_lokr_forward( GGML_ASSERT(k * k * h_split->ne[2] == a->ne[2]); a = ggml_reshape_4d(ctx, a, a->ne[0] * k, a->ne[1] * k, a->ne[2] / (k * k), a->ne[3]); } - ggml_tensor* ha = ggml_ext_conv_2d(ctx, h_split, a, nullptr, + ggml_tensor* ha = ggml_ext_conv_2d(runner_ctx, h_split, a, nullptr, conv_params.s0, conv_params.s1, conv_params.p0, @@ -167,8 +201,7 @@ ggml_tensor* ggml_ext_lokr_forward( conv_params.scale); // not supporting lora_mid here - hb = ggml_ext_conv_2d(ctx, - ha, + hb = ggml_ext_conv_2d(runner_ctx, ha, b, nullptr, 1, @@ -204,6 +237,6 @@ ggml_tensor* ggml_ext_lokr_forward( ggml_tensor* hc = ggml_transpose(ctx, hc_t); // ungroup ggml_tensor* out = ggml_reshape_4d(ctx, ggml_cont(ctx, hc), w_out, h_out, up * vp, batch); - return ggml_ext_scale(ctx, out, scale); + return ggml_ext_scale(runner_ctx, out, scale); } } diff --git a/src/model/adapter/lora_ops.h b/src/model/adapter/lora_ops.h index b32debfa69..309f513653 100644 --- a/src/model/adapter/lora_ops.h +++ b/src/model/adapter/lora_ops.h @@ -3,23 +3,32 @@ #include "core/ggml_runner.h" -ggml_tensor* ggml_ext_merge_lora(ggml_context* ctx, +// n-mode tensor-matrix product +// example: 2-mode product +// A: [ne03, k, ne01, ne00] +// B: k rows, m columns => [k, m] +// result is [ne03, m, ne01, ne00] +ggml_tensor* ggml_ext_mul_n_mode(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b, int mode = 0); + +// Kronecker product +// [ne03,ne02,ne01,ne00] x [ne13,ne12,ne11,ne10] => [ne03*ne13,ne02*ne12,ne01*ne11,ne00*ne10] +ggml_tensor* ggml_ext_kronecker(GGMLRunnerContext* runner_ctx, ggml_tensor* a, ggml_tensor* b); + +ggml_tensor* ggml_ext_merge_lora(GGMLRunnerContext* runner_ctx, ggml_tensor* lora_down, ggml_tensor* lora_up, ggml_tensor* lora_mid = nullptr); -ggml_tensor* ggml_ext_lokr_forward( - ggml_context* ctx, - ggml_backend_t backend, - ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] - ggml_tensor* w1, // Outer C (Full rank) - ggml_tensor* w1a, // Outer A (Low rank part 1) - ggml_tensor* w1b, // Outer B (Low rank part 2) - ggml_tensor* w2, // Inner BA (Full rank) - ggml_tensor* w2a, // Inner A (Low rank part 1) - ggml_tensor* w2b, // Inner B (Low rank part 2) - bool is_conv, - WeightAdapter::ForwardParams::conv2d_params_t conv_params, - float scale); +ggml_tensor* ggml_ext_lokr_forward(GGMLRunnerContext* runner_ctx, + ggml_tensor* h, // Input: [q, batch] or [W, H, q, batch] + ggml_tensor* w1, // Outer C (Full rank) + ggml_tensor* w1a, // Outer A (Low rank part 1) + ggml_tensor* w1b, // Outer B (Low rank part 2) + ggml_tensor* w2, // Inner BA (Full rank) + ggml_tensor* w2a, // Inner A (Low rank part 1) + ggml_tensor* w2b, // Inner B (Low rank part 2) + bool is_conv, + WeightAdapter::ForwardParams::conv2d_params_t conv_params, + float scale); #endif // __SD_MODEL_ADAPTER_LORA_OPS_H__ diff --git a/src/model/adapter/pmid.hpp b/src/model/adapter/pmid.hpp index 900db5cb18..163d2562bb 100644 --- a/src/model/adapter/pmid.hpp +++ b/src/model/adapter/pmid.hpp @@ -36,7 +36,7 @@ struct FuseBlock : public GGMLBlock { x = layer_norm->forward(ctx, x); // x = ggml_add(ctx, ggml_mul_mat(ctx, fc1_w, x), fc1_b); x = fc1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = fc2->forward(ctx, x); // x = ggml_add(ctx, ggml_mul_mat(ctx, fc2_w, x), fc2_b); if (use_residue) @@ -132,8 +132,8 @@ struct PerceiverAttention : public GGMLBlock { k = reshape_tensor(ctx->ggml_ctx, k, heads); v = reshape_tensor(ctx->ggml_ctx, v, heads); scale = 1.f / sqrt(sqrt((float)dim_head)); - k = ggml_ext_scale(ctx->ggml_ctx, k, scale, true); - q = ggml_ext_scale(ctx->ggml_ctx, q, scale, true); + k = ggml_ext_scale(ctx, k, scale, true); + q = ggml_ext_scale(ctx, q, scale, true); // auto weight = ggml_mul_mat(ctx, q, k); auto weight = ggml_mul_mat(ctx->ggml_ctx, k, q); // NOTE order of mul is opposite to pytorch diff --git a/src/model/audio/wav2vec2.hpp b/src/model/audio/wav2vec2.hpp index 3f3c969ef1..744df47201 100644 --- a/src/model/audio/wav2vec2.hpp +++ b/src/model/audio/wav2vec2.hpp @@ -67,7 +67,7 @@ namespace Wav2Vec2 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { auto conv = std::dynamic_pointer_cast(blocks["conv"]); x = conv->forward(ctx, x); - return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x)); } }; @@ -85,7 +85,7 @@ namespace Wav2Vec2 { x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); x = layer_norm->forward(ctx, x); x = ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3); - return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x)); } }; @@ -103,7 +103,7 @@ namespace Wav2Vec2 { x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]); x = layer_norm->forward(ctx, x); x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[2], x->ne[3]); - return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, x)); + return ggml_gelu_erf_inplace(ctx->ggml_ctx, ggml_ext_cont(ctx, x)); } }; @@ -198,7 +198,7 @@ namespace Wav2Vec2 { auto w = weight(ctx); auto b = params.count("conv.bias") > 0 ? params["conv.bias"] : nullptr; x = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); - x = ggml_ext_conv_1d(ctx->ggml_ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true); + x = ggml_ext_conv_1d(ctx, x, w, b, 1, kernel_size_ / 2, 1, groups_, true); // Apply GELU out of place before cropping to keep graph buffer reuse safe. x = ggml_gelu_erf(ctx->ggml_ctx, x); x = ggml_view_3d(ctx->ggml_ctx, x, x->ne[0] - 1, x->ne[1], x->ne[2], x->nb[1], x->nb[2], 0); @@ -216,7 +216,7 @@ namespace Wav2Vec2 { auto intermediate_dense = std::dynamic_pointer_cast(blocks["intermediate_dense"]); auto output_dense = std::dynamic_pointer_cast(blocks["output_dense"]); x = intermediate_dense->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = output_dense->forward(ctx, x); return x; } diff --git a/src/model/common/block.hpp b/src/model/common/block.hpp index 5be333d6fa..bed820c240 100644 --- a/src/model/common/block.hpp +++ b/src/model/common/block.hpp @@ -32,7 +32,7 @@ class DownSampleBlock : public GGMLBlock { if (vae_downsample) { auto conv = std::dynamic_pointer_cast(blocks["conv"]); - x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); x = conv->forward(ctx, x); } else { auto conv = std::dynamic_pointer_cast(blocks["op"]); @@ -220,7 +220,7 @@ class GELU : public UnaryBlock { auto proj = std::dynamic_pointer_cast(blocks["proj"]); x = proj->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); return x; } }; @@ -250,7 +250,7 @@ struct Mlp : public GGMLBlock { auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); x = fc1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = fc2->forward(ctx, x); return x; } @@ -293,9 +293,7 @@ class FeedForward : public GGMLBlock { auto net_0 = std::dynamic_pointer_cast(blocks["net.0"]); auto net_2 = std::dynamic_pointer_cast(blocks["net.2"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - net_2->set_force_prec_f32(true); - } + net_2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); x = net_0->forward(ctx, x); // [ne3, ne2, ne1, inner_dim] x = net_2->forward(ctx, x); // [ne3, ne2, ne1, dim_out] @@ -597,7 +595,7 @@ class ResidualAttentionBlock : public GGMLBlock { x = ggml_add(ctx->ggml_ctx, x, attn->forward(ctx, ln_1->forward(ctx, x))); auto h = mlp_c_fc->forward(ctx, ln_2->forward(ctx, x)); - h = ggml_ext_gelu_quick(ctx->ggml_ctx, h); + h = ggml_ext_gelu_quick(ctx, h); h = mlp_c_proj->forward(ctx, h); return ggml_add(ctx->ggml_ctx, x, h); } @@ -614,7 +612,7 @@ class AlphaBlender : public GGMLBlock { ggml_tensor* get_alpha(GGMLRunnerContext* ctx) { // image_only_indicator is always tensor([0.]) and since mix_factor.shape is [1,] // so learned_with_images is same as learned - auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]); + auto mix_factor = ggml_ext_cast_f32(ctx, params["mix_factor"]); return ggml_sigmoid(ctx->ggml_ctx, mix_factor); } diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index 39172bc3f6..eb437e5df5 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -147,6 +147,7 @@ class Linear : public UnaryBlock { bool bias; bool force_f32; bool force_prec_f32; + uint32_t precision_flags = 0; bool has_weight_scale = false; int w4_convrot_kind = 0; bool int8_convrot = false; @@ -215,6 +216,10 @@ class Linear : public UnaryBlock { scale = scale_; } + void set_precision_flags(uint32_t flags) { + precision_flags = flags; + } + void set_force_prec_f32(bool force_prec_f32_) { force_prec_f32 = force_prec_f32_; } @@ -225,72 +230,34 @@ class Linear : public UnaryBlock { } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + const bool force_prec_f32 = ggml_ext_force_prec_f32(ctx, this->force_prec_f32, precision_flags); ggml_tensor* w = params["weight"]; const float scale = ctx->linear_scale > 0.f ? ctx->linear_scale : this->scale; ggml_tensor* weight_scale = has_weight_scale ? params["weight_scale"] : nullptr; -#ifndef SD_USE_UPSTREAM_GGML - if (w->type == GGML_TYPE_F8_E4M3 || w->type == GGML_TYPE_F8_E5M2) { - bool supports_fp8_matmul = false; - if (ctx->backend != nullptr) { - ggml_tensor* fp8_matmul = ggml_mul_mat(ctx->ggml_ctx, w, x); - if (force_prec_f32) { - ggml_prec_set_acc(fp8_matmul, GGML_PREC_F32); - } - supports_fp8_matmul = ggml_backend_supports_op(ctx->backend, fp8_matmul) && !sd_backend_is_cpu(ctx->backend); - } - if (!supports_fp8_matmul) { - w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_BF16); - } - } -#endif - ggml_tensor* b = nullptr; + w = ggml_ext_prepare_linear_weight(ctx, w, x, force_prec_f32); + ggml_tensor* b = nullptr; if (bias) { b = params["bias"]; } ggml_tensor* linear_bias = has_weight_scale ? nullptr : b; ggml_tensor* out = nullptr; if (w->type == GGML_TYPE_I8) { - if (x->type != GGML_TYPE_F32) { - x = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, x); - } - if (!ggml_is_contiguous(x)) { - x = ggml_cont(ctx->ggml_ctx, x); - } + x = ggml_ext_prepare_i8_input(ctx, x); ggml_tensor* lora_input = x; if (ctx->weight_adapter && b != nullptr) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); - } -#ifndef SD_USE_UPSTREAM_GGML - if (int8_convrot && scale == 1.f) { - const auto cache_key = std::make_pair(x, int8_convrot_group_size); - auto cached = ctx->int8_convrot_cache.find(cache_key); - if (cached == ctx->int8_convrot_cache.end()) { - x = ggml_quantize_i8_convrot(ctx->ggml_ctx, x, int8_convrot_group_size); - ctx->int8_convrot_cache.emplace(cache_key, x); - } else { - x = cached->second; - } + b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias"); } -#endif + #ifndef SD_USE_UPSTREAM_GGML if (w4_convrot_kind) { - if (scale != 1.f) - x = ggml_ext_scale(ctx->ggml_ctx, x, scale); - out = ggml_mul_mat_w4_convrot(ctx->ggml_ctx, w, x, - w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale, - w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr, - w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr, - scale == 1.f ? b : nullptr, static_cast(w4_convrot_kind), - int8_convrot ? int8_convrot_group_size : 0); - GGML_ASSERT(out); - if (scale != 1.f) { - out = ggml_ext_scale(ctx->ggml_ctx, out, 1.f / scale); - if (b) - out = ggml_add_inplace(ctx->ggml_ctx, out, b); - } + out = ggml_ext_linear_w4_convrot(ctx, x, w, + w4_convrot_kind == 2 ? params["weight_codebook"] : weight_scale, + w4_convrot_kind == 2 ? params["weight_s_channel"] : nullptr, + w4_convrot_kind == 2 ? params["weight_s_rel"] : nullptr, + b, static_cast(w4_convrot_kind), int8_convrot ? int8_convrot_group_size : 0, scale); } else #endif - out = ggml_ext_linear_i8_tensorwise(ctx->ggml_ctx, + out = ggml_ext_linear_i8_tensorwise(ctx, x, w, weight_scale, @@ -302,33 +269,21 @@ class Linear : public UnaryBlock { forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; forward_params.linear.force_prec_f32 = force_prec_f32; forward_params.linear.scale = scale; - out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, - ctx->backend, - lora_input, - w, - out, - prefix, - forward_params); + out = ctx->weight_adapter->add_lora_to_output(ctx, lora_input, w, out, prefix, forward_params); } return out; } if (has_weight_scale) { - out = ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr, force_prec_f32, scale); + out = ggml_ext_linear(ctx, x, w, nullptr, force_prec_f32, scale); out = ggml_mul(ctx->ggml_ctx, out, weight_scale); if (ctx->weight_adapter) { WeightAdapter::ForwardParams forward_params; forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; forward_params.linear.force_prec_f32 = force_prec_f32; forward_params.linear.scale = scale; - out = ctx->weight_adapter->add_lora_to_output(ctx->ggml_ctx, - ctx->backend, - x, - w, - out, - prefix, - forward_params); + out = ctx->weight_adapter->add_lora_to_output(ctx, x, w, out, prefix, forward_params); if (b != nullptr) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias"); } } if (b != nullptr) { @@ -341,9 +296,9 @@ class Linear : public UnaryBlock { forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_LINEAR; forward_params.linear.force_prec_f32 = force_prec_f32; forward_params.linear.scale = scale; - out = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, linear_bias, prefix, forward_params); + out = ctx->weight_adapter->forward_with_lora(ctx, x, w, linear_bias, prefix, forward_params); } else { - out = ggml_ext_linear(ctx->ggml_ctx, x, w, linear_bias, force_prec_f32, scale); + out = ggml_ext_linear(ctx, x, w, linear_bias, force_prec_f32, scale); } return out; } @@ -376,42 +331,6 @@ class Embedding : public UnaryBlock { bool is_int8_tensorwise = false; int int8_convrot_group_size = 0; - static void get_rows_i8(ggml_tensor* dst, int ith, int nth, void* userdata) { - const auto* embedding = static_cast(userdata); - const int group_size = embedding->int8_convrot_group_size; - const auto* weight = dst->src[0]; - const auto* input_ids = static_cast(dst->src[1]->data); - const auto* scales = static_cast(dst->src[2]->data); - const bool scalar_scale = ggml_nelements(dst->src[2]) == 1; - const float normalization = group_size > 0 ? 1.f / std::sqrt(static_cast(group_size)) : 1.f; - for (int64_t row = ith; row < dst->ne[1]; row += nth) { - const int32_t token = input_ids[row]; - GGML_ASSERT(token >= 0 && token < weight->ne[1]); - const auto* src = reinterpret_cast(static_cast(weight->data) + token * weight->nb[1]); - auto* out = reinterpret_cast(static_cast(dst->data) + row * dst->nb[1]); - const float scale = scales[scalar_scale ? 0 : token] * normalization; - for (int64_t i = 0; i < weight->ne[0]; ++i) { - out[i] = static_cast(src[i]) * scale; - } - // The regular Hadamard rotation is symmetric and its own inverse. - for (int stride = 1; stride < group_size; stride *= 4) { - for (int64_t base = 0; base < weight->ne[0]; base += 4 * stride) { - for (int j = 0; j < stride; ++j) { - float* values = out + base + j; - const float a = values[0]; - const float b = values[stride]; - const float c = values[2 * stride]; - const float d = values[3 * stride]; - values[0] = a + b + c - d; - values[stride] = a + b - c + d; - values[2 * stride] = a - b + c + d; - values[3 * stride] = -a + b + c + d; - } - } - } - } - } - void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map, const std::string prefix = "") override { auto weight_storage = tensor_storage_map.find(prefix + "weight"); is_int8_tensorwise = weight_storage != tensor_storage_map.end() && weight_storage->second.is_int8_tensorwise; @@ -451,30 +370,9 @@ class Embedding : public UnaryBlock { num_embeddings(num_embeddings) { } - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* input_ids) override { - // input_ids: [N, n_token] - auto weight = params["weight"]; - - // There are issues with ggml batch inference, so we are expanding it here first. - // TODO: fix ggml batch inference - int64_t n = input_ids->ne[1]; - input_ids = ggml_reshape_1d(ctx->ggml_ctx, input_ids, input_ids->ne[0] * input_ids->ne[1]); - - ggml_tensor* embedding; - if (is_int8_tensorwise) { - GGML_ASSERT(input_ids->type == GGML_TYPE_I32); - ggml_tensor* args[] = {weight, input_ids, params["weight_scale"]}; - embedding = ggml_custom_4d(ctx->ggml_ctx, GGML_TYPE_F32, embedding_dim, input_ids->ne[0], 1, 1, - args, 3, get_rows_i8, GGML_N_TASKS_MAX, this); - } else { - input_ids = ggml_reshape_3d(ctx->ggml_ctx, input_ids, input_ids->ne[0], 1, input_ids->ne[1]); - embedding = ggml_get_rows(ctx->ggml_ctx, weight, input_ids); - } - embedding = ggml_reshape_3d(ctx->ggml_ctx, embedding, embedding->ne[0], embedding->ne[1] / n, n); - - // [N, n_token, embedding_dim] - return embedding; + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input_ids) override { + return ggml_ext_embedding(ctx, params["weight"], input_ids, + is_int8_tensorwise ? params["weight_scale"] : nullptr, int8_convrot_group_size); } }; @@ -528,7 +426,7 @@ class Conv1d : public UnaryBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { GGML_ASSERT(x->ne[1] == in_channels); - return ggml_ext_conv_1d(ctx->ggml_ctx, x, params["weight"], bias ? params["bias"] : nullptr, + return ggml_ext_conv_1d(ctx, x, params["weight"], bias ? params["bias"] : nullptr, stride, padding, dilation, groups, force_prec_f32); } }; @@ -598,9 +496,9 @@ class Conv2d : public UnaryBlock { forward_params.conv2d.circular_x = ctx->circular_x_enabled; forward_params.conv2d.circular_y = ctx->circular_y_enabled; forward_params.conv2d.scale = scale; - return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); + return ctx->weight_adapter->forward_with_lora(ctx, x, w, b, prefix, forward_params); } - return ggml_ext_conv_2d(ctx->ggml_ctx, + return ggml_ext_conv_2d(ctx, x, w, b, @@ -687,9 +585,9 @@ class Conv2d_grouped : public UnaryBlock { forward_params.conv2d.circular_x = ctx->circular_x_enabled; forward_params.conv2d.circular_y = ctx->circular_y_enabled; forward_params.conv2d.scale = scale; - return ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x, w, b, prefix, forward_params); + return ctx->weight_adapter->forward_with_lora(ctx, x, w, b, prefix, forward_params); } - return ggml_ext_conv_2d(ctx->ggml_ctx, x, w, b, + return ggml_ext_conv_2d(ctx, x, w, b, stride.second, stride.first, padding.second, padding.first, dilation.second, dilation.first, @@ -699,24 +597,15 @@ class Conv2d_grouped : public UnaryBlock { scale); } - if (groups == in_channels && groups == out_channels) { - ggml_tensor* res; - if (ctx->conv2d_direct_enabled) { - res = ggml_conv_2d_dw_direct(ctx->ggml_ctx, w, x, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first); - } else { - res = ggml_conv_2d_dw(ctx->ggml_ctx, w, x, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first); - } - if (b) { - b = ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, b->ne[0], 1); - res = ggml_add_inplace(ctx->ggml_ctx, res, b); - } - return res; + if (!ctx->weight_adapter || (groups == in_channels && groups == out_channels)) { + return ggml_ext_conv_2d_grouped(ctx, x, w, b, groups, + stride.second, stride.first, + padding.second, padding.first, + dilation.second, dilation.first, + ctx->conv2d_direct_enabled, + ctx->circular_x_enabled, + ctx->circular_y_enabled, + scale); } int64_t ic_g = in_channels / groups; @@ -743,33 +632,22 @@ class Conv2d_grouped : public UnaryBlock { b_i = ggml_view_1d(ctx->ggml_ctx, b, oc_g, b_offset); } - if (ctx->weight_adapter) { - WeightAdapter::ForwardParams forward_params; - forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; - forward_params.conv2d.s0 = stride.second; - forward_params.conv2d.s1 = stride.first; - forward_params.conv2d.p0 = padding.second; - forward_params.conv2d.p1 = padding.first; - forward_params.conv2d.d0 = dilation.second; - forward_params.conv2d.d1 = dilation.first; - forward_params.conv2d.direct = ctx->conv2d_direct_enabled; - forward_params.conv2d.circular_x = ctx->circular_x_enabled; - forward_params.conv2d.circular_y = ctx->circular_y_enabled; - forward_params.conv2d.scale = scale; - out_slices[i] = ctx->weight_adapter->forward_with_lora(ctx->ggml_ctx, ctx->backend, x_i, w_i, b_i, prefix, forward_params); - } else { - out_slices[i] = ggml_ext_conv_2d(ctx->ggml_ctx, x_i, w_i, b_i, - stride.second, stride.first, - padding.second, padding.first, - dilation.second, dilation.first, - ctx->conv2d_direct_enabled, - ctx->circular_x_enabled, - ctx->circular_y_enabled, - scale); - } + WeightAdapter::ForwardParams forward_params; + forward_params.op_type = WeightAdapter::ForwardParams::op_type_t::OP_CONV2D; + forward_params.conv2d.s0 = stride.second; + forward_params.conv2d.s1 = stride.first; + forward_params.conv2d.p0 = padding.second; + forward_params.conv2d.p1 = padding.first; + forward_params.conv2d.d0 = dilation.second; + forward_params.conv2d.d1 = dilation.first; + forward_params.conv2d.direct = ctx->conv2d_direct_enabled; + forward_params.conv2d.circular_x = ctx->circular_x_enabled; + forward_params.conv2d.circular_y = ctx->circular_y_enabled; + forward_params.conv2d.scale = scale; + out_slices[i] = ctx->weight_adapter->forward_with_lora(ctx, x_i, w_i, b_i, prefix, forward_params); } - ggml_tensor* out = ggml_ext_vec_concat(ctx->ggml_ctx, out_slices, 2); + ggml_tensor* out = ggml_ext_vec_concat(ctx, out_slices, 2); return out; } @@ -823,7 +701,7 @@ class Conv3d : public UnaryBlock { ggml_tensor* w = params["weight"]; ggml_tensor* b = nullptr; if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight"); if (w->type != GGML_TYPE_F16) { w = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F16); } @@ -831,10 +709,10 @@ class Conv3d : public UnaryBlock { if (bias) { b = params["bias"]; if (ctx->weight_adapter) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias"); } } - return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, + return ggml_ext_conv_3d(ctx, x, w, b, in_channels, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), std::get<2>(padding), std::get<1>(padding), std::get<0>(padding), std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), @@ -879,16 +757,16 @@ class LayerNorm : public UnaryBlock { if (elementwise_affine) { w = params["weight"]; if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight"); } if (bias) { b = params["bias"]; if (ctx->weight_adapter) { - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias"); } } } - return ggml_ext_layer_norm(ctx->ggml_ctx, x, w, b, eps); + return ggml_ext_layer_norm(ctx, x, w, b, eps); } }; @@ -927,11 +805,11 @@ class GroupNorm : public GGMLBlock { w = params["weight"]; b = params["bias"]; if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); - b = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, b, prefix + "bias"); + w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight"); + b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias"); } } - return ggml_ext_group_norm(ctx->ggml_ctx, x, w, b, num_groups, eps); + return ggml_ext_group_norm(ctx, x, w, b, num_groups, eps); } }; @@ -971,7 +849,7 @@ class RMSNorm : public UnaryBlock { } ggml_tensor* w = params["weight"]; if (ctx->weight_adapter) { - w = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, w, prefix + "weight"); + w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight"); } x = ggml_rms_norm(ctx->ggml_ctx, x, eps); // The affine multiply may run on a different backend after normalization falls back. @@ -1032,7 +910,7 @@ class MultiheadAttention : public GGMLBlock { if (proj_in) { auto in_proj = std::dynamic_pointer_cast(blocks[in_proj_name]); auto qkv = in_proj->forward(ctx, x); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); q = qkv_vec[0]; k = qkv_vec[1]; v = qkv_vec[2]; diff --git a/src/model/common/rope.hpp b/src/model/common/rope.hpp index b6fdd328e1..b6f3111a09 100644 --- a/src/model/common/rope.hpp +++ b/src/model/common/rope.hpp @@ -1071,58 +1071,7 @@ namespace Rope { ggml_tensor* x, ggml_tensor* pe, bool rope_interleaved = true) { - auto ctx = runner_ctx->ggml_ctx; -#ifndef SD_USE_UPSTREAM_GGML - auto backend = runner_ctx->backend; - 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 = sd_ggml_rope_apply(ctx, backend, x, pe, rope_interleaved); - if (out && ggml_backend_supports_op(backend, out)) { - return out; - } - } -#endif - // x: [N, L, n_head, d_head] - // pe: [L, d_head/2, 2, 2], [[cos, -sin], [sin, cos]] - int64_t d_head = x->ne[0]; - int64_t n_head = x->ne[1]; - int64_t L = x->ne[2]; - int64_t N = x->ne[3]; - x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N, n_head, L, d_head] - if (rope_interleaved) { - x = ggml_reshape_4d(ctx, x, 2, d_head / 2, L, n_head * N); // [N * n_head, L, d_head/2, 2] - x = ggml_cont(ctx, ggml_permute(ctx, x, 3, 0, 1, 2)); // [2, N * n_head, L, d_head/2] - } else { - auto x_0 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], 0); - auto x_1 = ggml_view_3d(ctx, x, d_head / 2, L, n_head * N, x->nb[1], x->nb[2], (d_head / 2) * x->nb[0]); - - auto pe_2 = ggml_cont(ctx, ggml_permute(ctx, pe, 2, 3, 0, 1)); - auto cos = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], 0); - auto sin = ggml_view_2d(ctx, pe_2, d_head / 2, L, pe_2->nb[1], pe_2->nb[3]); - - auto part_0 = ggml_sub(ctx, ggml_mul(ctx, x_0, cos), ggml_mul(ctx, x_1, sin)); - auto part_1 = ggml_add(ctx, ggml_mul(ctx, x_0, sin), ggml_mul(ctx, x_1, cos)); - return ggml_concat(ctx, part_0, part_1, 0); - } - - int64_t offset = x->nb[2] * x->ne[2]; - auto x_0 = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 0); // [N * n_head, L, d_head/2] - auto x_1 = ggml_view_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2], x->nb[1], x->nb[2], offset * 1); // [N * n_head, L, d_head/2] - x_0 = ggml_reshape_4d(ctx, x_0, 1, x_0->ne[0], x_0->ne[1], x_0->ne[2]); // [N * n_head, L, d_head/2, 1] - x_1 = ggml_reshape_4d(ctx, x_1, 1, x_1->ne[0], x_1->ne[1], x_1->ne[2]); // [N * n_head, L, d_head/2, 1] - auto temp_x = ggml_new_tensor_4d(ctx, x_0->type, 2, x_0->ne[1], x_0->ne[2], x_0->ne[3]); - x_0 = ggml_repeat(ctx, x_0, temp_x); // [N * n_head, L, d_head/2, 2] - x_1 = ggml_repeat(ctx, x_1, temp_x); // [N * n_head, L, d_head/2, 2] - - pe = ggml_cont(ctx, ggml_permute(ctx, pe, 3, 0, 1, 2)); // [2, L, d_head/2, 2] - offset = pe->nb[2] * pe->ne[2]; - auto pe_0 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 0); // [L, d_head/2, 2] - auto pe_1 = ggml_view_3d(ctx, pe, pe->ne[0], pe->ne[1], pe->ne[2], pe->nb[1], pe->nb[2], offset * 1); // [L, d_head/2, 2] - - auto x_out = ggml_add_inplace(ctx, ggml_mul(ctx, x_0, pe_0), ggml_mul(ctx, x_1, pe_1)); // [N * n_head, L, d_head/2, 2] - x_out = ggml_reshape_3d(ctx, x_out, d_head, L, n_head * N); // [N*n_head, L, d_head] - return x_out; + return ggml_ext_rope_apply(runner_ctx, x, pe, rope_interleaved); } __STATIC_INLINE__ ggml_tensor* attention(GGMLRunnerContext* ctx, diff --git a/src/model/common/vsr.hpp b/src/model/common/vsr.hpp index c203dff10a..4d4fa7a263 100644 --- a/src/model/common/vsr.hpp +++ b/src/model/common/vsr.hpp @@ -177,7 +177,7 @@ namespace VSR { params.window_height[i] = windows[i].first; params.window_width[i] = windows[i].second; } - if (auto out = sd_ggml_window_attn(c, ctx->backend, q, k, v, params)) + if (auto out = ggml_ext_window_attn(ctx, q, k, v, params)) return out; } #endif @@ -192,7 +192,7 @@ namespace VSR { auto tk = ggml_get_rows(c, k, ki); auto tv = ggml_get_rows(c, v, ki); auto mask = inputs.input(ctx, tile.mask); - auto out = ggml_ext_attention_ext(c, ctx->backend, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled); + auto out = ggml_ext_attention_ext(ctx, tq, tk, tv, heads, mask, false, ctx->flash_attn_enabled); out = ggml_reshape_3d(c, out, out->ne[0], tile.width, tile.height); row = row ? ggml_concat(c, row, out, 1) : out; row_width += tile.width; diff --git a/src/model/diffusion/anima.hpp b/src/model/diffusion/anima.hpp index e9e33b88a1..08c20fe683 100644 --- a/src/model/diffusion/anima.hpp +++ b/src/model/diffusion/anima.hpp @@ -123,7 +123,7 @@ namespace Anima { emb = ggml_add(ctx->ggml_ctx, emb, temb); } - auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 3, 0); + auto emb_chunks = ggml_ext_chunk(ctx, emb, 3, 0); auto shift = emb_chunks[0]; auto scale = emb_chunks[1]; auto gate = emb_chunks[2]; @@ -164,7 +164,7 @@ namespace Anima { emb = ggml_add(ctx->ggml_ctx, emb, temb_2c); } - auto emb_chunks = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0); + auto emb_chunks = ggml_ext_chunk(ctx, emb, 2, 0); auto shift = emb_chunks[0]; auto scale = emb_chunks[1]; @@ -275,7 +275,7 @@ namespace Anima { auto layer2 = std::dynamic_pointer_cast(blocks["layer2"]); x = layer1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = layer2->forward(ctx, x); return x; } @@ -293,7 +293,7 @@ namespace Anima { auto layer2 = std::dynamic_pointer_cast(blocks["2"]); x = layer0->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = layer2->forward(ctx, x); return x; } @@ -500,14 +500,14 @@ namespace Anima { int64_t W = x->ne[0]; int64_t H = x->ne[1]; - auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, x->ne[0], x->ne[1], 1, x->ne[3]); + auto padding_mask = ggml_ext_zeros(ctx, x->ne[0], x->ne[1], 1, x->ne[3]); x = ggml_concat(ctx->ggml_ctx, x, padding_mask, 2); // [N, C + 1, H, W] x = DiT::pad_and_patchify(ctx, x, config.patch_size, config.patch_size); // [N, h*w, (C+1)*ph*pw] int64_t img_len = x->ne[1]; if (ref_latents.size() > 0) { for (ggml_tensor* ref : ref_latents) { - auto padding_mask = ggml_ext_zeros(ctx->ggml_ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]); + auto padding_mask = ggml_ext_zeros(ctx, ref->ne[0], ref->ne[1], 1, ref->ne[3]); ref = ggml_concat(ctx->ggml_ctx, ref, padding_mask, 2); // [N, C + 1, H, W] ref = DiT::pad_and_patchify(ctx, ref, config.patch_size, config.patch_size); x = ggml_concat(ctx->ggml_ctx, x, ref, 1); @@ -515,7 +515,7 @@ namespace Anima { } x = x_embedder->forward(ctx, x); - auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(config.hidden_size)); + auto timestep_proj = ggml_ext_timestep_embedding(ctx, timestep, static_cast(config.hidden_size)); auto temb = t_embedder->forward(ctx, timestep_proj); auto embedded_timestep = t_embedding_norm->forward(ctx, timestep_proj); @@ -530,14 +530,14 @@ namespace Anima { adapted_context = ggml_mul(ctx->ggml_ctx, adapted_context, w); } if (adapted_context->ne[1] < 512) { - auto pad_ctx = ggml_ext_zeros(ctx->ggml_ctx, + auto pad_ctx = ggml_ext_zeros(ctx, adapted_context->ne[0], 512 - adapted_context->ne[1], adapted_context->ne[2], 1); adapted_context = ggml_concat(ctx->ggml_ctx, adapted_context, pad_ctx, 1); } else if (adapted_context->ne[1] > 512) { - adapted_context = ggml_ext_slice(ctx->ggml_ctx, adapted_context, 1, 0, 512); + adapted_context = ggml_ext_slice(ctx, adapted_context, 1, 0, 512); } encoder_hidden_states = adapted_context; } @@ -552,11 +552,11 @@ namespace Anima { x = block->forward(ctx, x, encoder_hidden_states, embedded_timestep, temb, image_pe); sd::ggml_graph_cut::mark_graph_cut(x, "anima.blocks." + std::to_string(i), "x"); } - x = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, img_len); + x = ggml_ext_slice(ctx, x, 1, 0, img_len); x = final_layer->forward(ctx, x, embedded_timestep, temb); // [N, h*w, ph*pw*C] - x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, false); // [N, C, H, W] + x = DiT::unpatchify_and_crop(ctx, x, H, W, config.patch_size, config.patch_size, false); // [N, C, H, W] return x; } diff --git a/src/model/diffusion/animatediff.hpp b/src/model/diffusion/animatediff.hpp index eb0ffba31a..459e473b88 100644 --- a/src/model/diffusion/animatediff.hpp +++ b/src/model/diffusion/animatediff.hpp @@ -121,13 +121,13 @@ namespace AnimateDiff { auto residual = x; auto h = norm->forward(ctx, x); - h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1)); + h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1)); h = ggml_reshape_3d(ctx->ggml_ctx, h, C, num_frames, W * H); h = proj_in->forward(ctx, h); h = tb0->forward(ctx, h); h = proj_out->forward(ctx, h); h = ggml_reshape_4d(ctx->ggml_ctx, h, C, num_frames, W, H); - h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1)); + h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 2, 3, 0, 1)); return ggml_add(ctx->ggml_ctx, h, residual); } diff --git a/src/model/diffusion/boogu.hpp b/src/model/diffusion/boogu.hpp index ecdb7def3e..c10da05b98 100644 --- a/src/model/diffusion/boogu.hpp +++ b/src/model/diffusion/boogu.hpp @@ -161,7 +161,7 @@ namespace Boogu { auto caption_embedder_0 = std::dynamic_pointer_cast(blocks["caption_embedder.0"]); auto caption_embedder_1 = std::dynamic_pointer_cast(blocks["caption_embedder.1"]); - auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(frequency_embedding_size), 10000, timestep_scale); + auto timestep_proj = ggml_ext_timestep_embedding(ctx, timestep, static_cast(frequency_embedding_size), 10000, timestep_scale); auto time_embed = timestep_embedder->forward(ctx, timestep_proj); auto caption_embed = caption_embedder_1->forward(ctx, caption_embedder_0->forward(ctx, text_hidden_states)); return {time_embed, caption_embed}; @@ -179,7 +179,7 @@ namespace Boogu { auto norm = std::dynamic_pointer_cast(blocks["norm"]); emb = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, emb)); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 4, 0); + auto mods = ggml_ext_chunk(ctx, emb, 4, 0); auto scale_msa = mods[0]; auto gate_msa = mods[1]; @@ -204,9 +204,7 @@ namespace Boogu { auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); auto linear_3 = std::dynamic_pointer_cast(blocks["linear_3"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - linear_2->set_force_prec_f32(true); - } + linear_2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); auto h1 = linear_1->forward(ctx, x); auto h2 = linear_3->forward(ctx, x); @@ -264,9 +262,7 @@ namespace Boogu { auto norm_k = std::dynamic_pointer_cast(blocks["norm_k"]); auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - to_out_0->set_force_prec_f32(true); - } + to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); int64_t N = hidden_states->ne[2]; int64_t Lq = hidden_states->ne[1]; @@ -388,9 +384,7 @@ namespace Boogu { auto instruct_out = std::dynamic_pointer_cast(blocks["processor.instruct_out"]); auto img_out = std::dynamic_pointer_cast(blocks["processor.img_out"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - to_out_0->set_force_prec_f32(true); - } + to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); int64_t N = img_hidden_states->ne[2]; int64_t L_img = img_hidden_states->ne[1]; @@ -417,8 +411,8 @@ namespace Boogu { k = norm_k->forward(ctx, k); auto hidden_states = Rope::attention(ctx, q, k, v, rotary_emb, attention_mask); - auto instruct_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, L_instruct); - auto img_attn = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, L_instruct, L_instruct + L_img); + auto instruct_attn = ggml_ext_slice(ctx, hidden_states, 1, 0, L_instruct); + auto img_attn = ggml_ext_slice(ctx, hidden_states, 1, L_instruct, L_instruct + L_img); instruct_attn = instruct_out->forward(ctx, instruct_attn); img_attn = img_out->forward(ctx, img_attn); @@ -504,8 +498,8 @@ namespace Boogu { auto instruct_shift_mlp = std::get<1>(instruct_norm2_out_vec); auto joint_attn_out = img_instruct_attn->forward(ctx, img_norm1_out, instruct_norm1_out, joint_rotary_emb); - auto instruct_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, 0, L_instruct); - auto img_attn_out = ggml_ext_slice(ctx->ggml_ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]); + auto instruct_attn_out = ggml_ext_slice(ctx, joint_attn_out, 1, 0, L_instruct); + auto img_attn_out = ggml_ext_slice(ctx, joint_attn_out, 1, L_instruct, joint_attn_out->ne[1]); auto img_self_attn_out = img_self_attn->forward(ctx, img_norm3_out, img_norm3_out, img_rotary_emb); @@ -634,7 +628,7 @@ namespace Boogu { auto time_caption_embed = std::dynamic_pointer_cast(blocks["time_caption_embed"]); auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); - auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, timesteps), timesteps); + auto timestep = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx, timesteps), timesteps); auto embeds = time_caption_embed->forward(ctx, timestep, context); auto temb = embeds.first; auto txt = embeds.second; @@ -647,8 +641,8 @@ namespace Boogu { int64_t txt_len = txt->ne[1]; GGML_ASSERT(pe->ne[3] == txt_len + ref_len + img_len); - auto txt_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, 0, txt_len); - auto noise_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len); + auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt_len); + auto noise_pe = ggml_ext_slice(ctx, pe, 3, txt_len + ref_len, txt_len + ref_len + img_len); for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); @@ -664,7 +658,7 @@ namespace Boogu { ggml_tensor* combined_img = img; if (ref_img != nullptr) { - auto ref_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + ref_len); + auto ref_pe = ggml_ext_slice(ctx, pe, 3, txt_len, txt_len + ref_len); for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["ref_image_refiner." + std::to_string(i)]); ref_img = block->forward(ctx, ref_img, ref_pe, temb); @@ -673,7 +667,7 @@ namespace Boogu { combined_img = ggml_concat(ctx->ggml_ctx, ref_img, img, 1); } - auto img_pe = ggml_ext_slice(ctx->ggml_ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]); + auto img_pe = ggml_ext_slice(ctx, pe, 3, txt_len, txt_len + combined_img->ne[1]); for (int i = 0; i < config.num_double_stream_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["double_stream_layers." + std::to_string(i)]); auto result = block->forward(ctx, combined_img, txt, pe, img_pe, temb); @@ -691,9 +685,9 @@ namespace Boogu { } hidden_states = norm_out->forward(ctx, hidden_states, temb); - hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]); - hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, false); - hidden_states = ggml_ext_scale(ctx->ggml_ctx, hidden_states, -1.f); + hidden_states = ggml_ext_slice(ctx, hidden_states, 1, hidden_states->ne[1] - img_len, hidden_states->ne[1]); + hidden_states = DiT::unpatchify_and_crop(ctx, hidden_states, H, W, config.patch_size, config.patch_size, false); + hidden_states = ggml_ext_scale(ctx, hidden_states, -1.f); return hidden_states; } }; diff --git a/src/model/diffusion/control.hpp b/src/model/diffusion/control.hpp index 7848892d0f..4ad85ecf8c 100644 --- a/src/model/diffusion/control.hpp +++ b/src/model/diffusion/control.hpp @@ -230,7 +230,7 @@ class ControlNetBlock : public GGMLBlock { auto linear_2 = std::dynamic_pointer_cast(blocks["control_add_embedding.linear_2"]); // [N*num_control_types, CONTROL_TYPE_EMBED_DIM] - auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM); + auto e = ggml_ext_timestep_embedding(ctx, control_types, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM); e = ggml_reshape_2d(ctx->ggml_ctx, e, ControlNetUnionConfig::CONTROL_TYPE_EMBED_DIM * config.num_control_types, batch); e = linear_1->forward(ctx, e); e = ggml_silu_inplace(ctx->ggml_ctx, e); @@ -238,8 +238,9 @@ class ControlNetBlock : public GGMLBlock { } // [N, C, H, W] => [N, C] - static ggml_tensor* spatial_mean(ggml_context* ctx, ggml_tensor* x) { - auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); + static ggml_tensor* spatial_mean(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + auto ctx = runner_ctx->ggml_ctx; + auto flat = ggml_reshape_3d(ctx, ggml_ext_cont(runner_ctx, x), x->ne[0] * x->ne[1], x->ne[2], x->ne[3]); return ggml_reshape_2d(ctx, ggml_mean(ctx, flat), x->ne[2], x->ne[3]); } @@ -255,11 +256,11 @@ class ControlNetBlock : public GGMLBlock { auto task_embedding = params["task_embedding"]; // [num_control_types, C] auto task = ggml_view_2d(c, task_embedding, channels, 1, task_embedding->nb[1], control_type_index * task_embedding->nb[1]); - auto hint_feat = ggml_add(c, spatial_mean(c, hint), task); // [N, C] + auto hint_feat = ggml_add(c, spatial_mean(ctx, hint), task); // [N, C] if (hint_feat->ne[1] != batch) { hint_feat = ggml_repeat(c, hint_feat, ggml_new_tensor_2d(c, GGML_TYPE_F32, channels, batch)); } - auto sample_feat = spatial_mean(c, sample); // [N, C] + auto sample_feat = spatial_mean(ctx, sample); // [N, C] // as in the reference (nn.MultiheadAttention without batch_first) the attention is over // the batch, and the 2 features (hint, sample) are its batch: [2, N, C] @@ -273,7 +274,7 @@ class ControlNetBlock : public GGMLBlock { } auto projection = std::dynamic_pointer_cast(blocks["spatial_ch_projs"]); - auto alpha = ggml_reshape_2d(c, ggml_ext_slice(c, x, 2, 0, 1), channels, batch); + auto alpha = ggml_reshape_2d(c, ggml_ext_slice(ctx, x, 2, 0, 1), channels, batch); alpha = projection->forward(ctx, alpha); alpha = ggml_reshape_4d(c, alpha, 1, 1, channels, batch); @@ -346,7 +347,7 @@ class ControlNetBlock : public GGMLBlock { auto middle_block_out = std::dynamic_pointer_cast(blocks["middle_block_out.0"]); - auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels); // [N, model_channels] + auto t_emb = ggml_ext_timestep_embedding(ctx, timesteps, model_channels); // [N, model_channels] auto emb = time_embed_0->forward(ctx, t_emb); emb = ggml_silu_inplace(ctx->ggml_ctx, emb); diff --git a/src/model/diffusion/dit.hpp b/src/model/diffusion/dit.hpp index c709b8ab4a..f82ad28532 100644 --- a/src/model/diffusion/dit.hpp +++ b/src/model/diffusion/dit.hpp @@ -5,11 +5,12 @@ #include "core/ggml_runner.h" namespace DiT { - inline ggml_tensor* patchify(ggml_context* ctx, + inline ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int pw, int ph, bool patch_last = true) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, C, H, W] // return: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C] int64_t N = x->ne[3]; @@ -28,19 +29,20 @@ namespace DiT { x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [N, h*w, C, ph*pw] x = ggml_reshape_3d(ctx, x, pw * ph * C, w * h, N); // [N, h*w, C*ph*pw] } else { - x = ggml_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, h*w, C, ph*pw] - x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N); // [N, h*w, ph*pw*C] + x = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3)); // [N, h*w, C, ph*pw] + x = ggml_reshape_3d(ctx, x, C * pw * ph, w * h, N); // [N, h*w, ph*pw*C] } return x; } - inline ggml_tensor* unpatchify(ggml_context* ctx, + inline ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t h, int64_t w, int ph, int pw, bool patch_last = true) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, h*w, C*ph*pw] if patch_last else [N, h*w, ph*pw*C] // return: [N, C, H, W] int64_t N = x->ne[2]; @@ -74,7 +76,7 @@ namespace DiT { int pad_h = (ph - H % ph) % ph; int pad_w = (pw - W % pw) % pw; - x = ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + x = ggml_ext_pad(ctx, x, pad_w, pad_h, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return x; } @@ -84,11 +86,11 @@ namespace DiT { int pw, bool patch_last = true) { x = pad_to_patch_size(ctx, x, ph, pw); - x = patchify(ctx->ggml_ctx, x, ph, pw, patch_last); + x = patchify(ctx, x, ph, pw, patch_last); return x; } - inline ggml_tensor* unpatchify_and_crop(ggml_context* ctx, + inline ggml_tensor* unpatchify_and_crop(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t H, int64_t W, @@ -99,19 +101,20 @@ namespace DiT { int pad_w = (pw - W % pw) % pw; int64_t h = ((H + pad_h) / ph); int64_t w = ((W + pad_w) / pw); - x = unpatchify(ctx, x, h, w, ph, pw, patch_last); // [N, C, H + pad_h, W + pad_w] - x = ggml_ext_slice(ctx, x, 1, 0, H); // [N, C, H, W + pad_w] - x = ggml_ext_slice(ctx, x, 0, 0, W); // [N, C, H, W] + x = unpatchify(runner_ctx, x, h, w, ph, pw, patch_last); // [N, C, H + pad_h, W + pad_w] + x = ggml_ext_slice(runner_ctx, x, 1, 0, H); // [N, C, H, W + pad_w] + x = ggml_ext_slice(runner_ctx, x, 0, 0, W); // [N, C, H, W] return x; } - inline ggml_tensor* patchify_3d(ggml_context* ctx, + inline ggml_tensor* patchify_3d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int pt, int ph, int pw, int64_t N = 1, bool patch_last = true) { + auto ctx = runner_ctx->ggml_ctx; // x: [N*C, T, H, W] // return: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C] int64_t C = x->ne[3] / N; @@ -125,24 +128,24 @@ namespace DiT { GGML_ASSERT(C * N == x->ne[3]); GGML_ASSERT(t_len * pt == T && h_len * ph == H && w_len * pw == W); - x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw] - x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt, t_len * C * N); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph, h_len * t_len * C * N); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw, w_len, ph * pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_reshape_4d(ctx, x, pw * ph * pt, w_len * h_len * t_len, C, N); // [N, C, t_len*h_len*w_len, pt*ph*pw] if (patch_last) { - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N, t_len*h_len*w_len, C, pt*ph*pw] } else { - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3)); // [N, t_len*h_len*w_len, pt*ph*pw, C] } // [N, t_len*h_len*w_len, C*pt*ph*pw] or [N, t_len*h_len*w_len, pt*ph*pw*C] x = ggml_reshape_4d(ctx, x, pw * ph * pt * C, w_len * h_len * t_len, N, 1); return x; } - inline ggml_tensor* unpatchify_3d(ggml_context* ctx, + inline ggml_tensor* unpatchify_3d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t t_len, int64_t h_len, @@ -151,6 +154,7 @@ namespace DiT { int ph, int pw, bool patch_last = true) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, t_len*h_len*w_len, C*pt*ph*pw] if patch_last else [N, t_len*h_len*w_len, pt*ph*pw*C] // return: [N*C, t_len*pt, h_len*ph, w_len*pw] int64_t N = x->ne[2]; @@ -159,20 +163,20 @@ namespace DiT { GGML_ASSERT(C * pt * ph * pw == x->ne[0]); if (patch_last) { - x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw * ph * pt, C, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, C, pt*ph*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] } else { - x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] } - x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] return x; } } // namespace DiT diff --git a/src/model/diffusion/ernie_image.hpp b/src/model/diffusion/ernie_image.hpp index 5898ac0679..532b8b846f 100644 --- a/src/model/diffusion/ernie_image.hpp +++ b/src/model/diffusion/ernie_image.hpp @@ -85,11 +85,12 @@ namespace ErnieImage { } }; - __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(GGMLRunnerContext* runner_ctx, ggml_tensor* timesteps, int dim, int max_period = 10000) { - auto emb = ggml_ext_timestep_embedding(ctx, timesteps, dim, max_period, 1.0f); + auto ctx = runner_ctx->ggml_ctx; + auto emb = ggml_ext_timestep_embedding(runner_ctx, timesteps, dim, max_period, 1.0f); int64_t half = dim / 2; auto cos_part = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], 0); auto sin_part = ggml_view_2d(ctx, emb, half, emb->ne[1], emb->nb[1], half * emb->nb[0]); @@ -97,7 +98,8 @@ namespace ErnieImage { return sin_first; } - __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(ggml_context* ctx, ggml_tensor* x, ggml_tensor* pe) { + __STATIC_INLINE__ ggml_tensor* apply_rotary_emb(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, S, heads, head_dim] // pe: [2, S, 1, head_dim], stored as ggml [head_dim, 1, S, 2]. int64_t head_dim = x->ne[0]; @@ -110,8 +112,8 @@ namespace ErnieImage { GGML_ASSERT(pe->ne[1] == 1 && pe->ne[2] == S && pe->ne[3] == 2); x = ggml_cont(ctx, x); - auto x_rot = ggml_ext_slice(ctx, x, 0, 0, rot_dim, false); - auto x_pass = rot_dim < head_dim ? ggml_ext_slice(ctx, x, 0, rot_dim, head_dim, false) : nullptr; + auto x_rot = ggml_ext_slice(runner_ctx, x, 0, 0, rot_dim, false); + auto x_pass = rot_dim < head_dim ? ggml_ext_slice(runner_ctx, x, 0, rot_dim, head_dim, false) : nullptr; int64_t half = rot_dim / 2; auto x1 = ggml_view_4d(ctx, x_rot, half, heads, S, N, x_rot->nb[1], x_rot->nb[2], x_rot->nb[3], 0); @@ -120,8 +122,8 @@ namespace ErnieImage { x2 = ggml_cont(ctx, x2); auto rotated = ggml_concat(ctx, ggml_neg(ctx, x2), x1, 0); - auto cos_emb = ggml_ext_slice(ctx, pe, 3, 0, 1, false); - auto sin_emb = ggml_ext_slice(ctx, pe, 3, 1, 2, false); + auto cos_emb = ggml_ext_slice(runner_ctx, pe, 3, 0, 1, false); + auto sin_emb = ggml_ext_slice(runner_ctx, pe, 3, 1, 2, false); auto out = ggml_add(ctx, ggml_mul(ctx, x_rot, cos_emb), ggml_mul(ctx, rotated, sin_emb)); if (x_pass != nullptr) { @@ -175,8 +177,8 @@ namespace ErnieImage { q = norm_q->forward(ctx, q); k = norm_k->forward(ctx, k); - q = apply_rotary_emb(ctx->ggml_ctx, q, pe); - k = apply_rotary_emb(ctx->ggml_ctx, k, pe); + q = apply_rotary_emb(ctx, q, pe); + k = apply_rotary_emb(ctx, k, pe); q = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)); // [N, heads, S, head_dim] q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]); @@ -269,7 +271,7 @@ namespace ErnieImage { auto norm = std::dynamic_pointer_cast(blocks["norm"]); auto linear = std::dynamic_pointer_cast(blocks["linear"]); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, conditioning), 2, 0); + auto mods = ggml_ext_chunk(ctx, linear->forward(ctx, conditioning), 2, 0); auto scale = mods[0]; auto shift = mods[1]; @@ -336,9 +338,9 @@ namespace ErnieImage { auto final_norm = std::dynamic_pointer_cast(blocks["final_norm"]); auto final_linear = std::dynamic_pointer_cast(blocks["final_linear"]); - auto img = x_embedder_proj->forward(ctx, x); // [N, hidden_size, Hp, Wp] - img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N); // [N, hidden_size, image_tokens] - img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); // [N, image_tokens, hidden_size] + auto img = x_embedder_proj->forward(ctx, x); // [N, hidden_size, Hp, Wp] + img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], N); // [N, hidden_size, image_tokens] + img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3)); // [N, image_tokens, hidden_size] auto txt = context; auto text_proj = std::dynamic_pointer_cast(blocks["text_proj"]); @@ -348,13 +350,13 @@ namespace ErnieImage { auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1); // [N, image_tokens + text_tokens, hidden_size] - auto sample = timestep_embedding_sin_cos(ctx->ggml_ctx, timestep, static_cast(config.hidden_size)); + auto sample = timestep_embedding_sin_cos(ctx, timestep, static_cast(config.hidden_size)); auto c = time_embedding->forward(ctx, sample); // [N, hidden_size] auto mod_params = adaLN_mod->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 6 * hidden_size] sd::ggml_graph_cut::mark_graph_cut(hidden_states, "ernie_image.prelude", "hidden_states"); // sd::ggml_graph_cut::mark_graph_cut(mod_params, "ernie_image.prelude", "mod_params"); - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, mod_params, 6, 0); + auto chunks = ggml_ext_chunk(ctx, mod_params, 6, 0); std::vector temb; temb.reserve(6); for (auto chunk : chunks) { @@ -368,10 +370,10 @@ namespace ErnieImage { } hidden_states = final_norm->forward(ctx, hidden_states, c); - hidden_states = final_linear->forward(ctx, hidden_states); // [N, image_tokens, p*p*out_channels] - auto patches = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, 0, n_img); // [N, image_tokens, hidden_size] + hidden_states = final_linear->forward(ctx, hidden_states); // [N, image_tokens, p*p*out_channels] + auto patches = ggml_ext_slice(ctx, hidden_states, 1, 0, n_img); // [N, image_tokens, hidden_size] - auto out = DiT::unpatchify(ctx->ggml_ctx, + auto out = DiT::unpatchify(ctx, patches, Hp, Wp, diff --git a/src/model/diffusion/fibo2.hpp b/src/model/diffusion/fibo2.hpp index 2a65ea1d9c..1ac7684320 100644 --- a/src/model/diffusion/fibo2.hpp +++ b/src/model/diffusion/fibo2.hpp @@ -108,14 +108,14 @@ namespace Fibo2 { }; ggml_tensor *q, *k, *v; if (fused_qkv) { - auto parts = ggml_ext_chunk(c, project("qkv", x), 3, 0); + auto parts = ggml_ext_chunk(ctx, 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); + auto parts = ggml_ext_chunk(ctx, project("to_kv", kv), 2, 0); k = parts[0]; v = parts[1]; } else { @@ -158,9 +158,7 @@ namespace Fibo2 { 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); - } + w2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); x = std::dynamic_pointer_cast(blocks["w13"])->forward(ctx, x); return w2->forward(ctx, ggml_swiglu(ctx->ggml_ctx, x)); } @@ -206,7 +204,7 @@ namespace Fibo2 { } 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); + auto x = ggml_ext_slice(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); } @@ -244,8 +242,8 @@ namespace Fibo2 { 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 mn = ggml_ext_chunk(ctx, modulation->forward(ctx, noisy), 4, 0); + auto mc = ggml_ext_chunk(ctx, 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]; @@ -311,12 +309,12 @@ namespace Fibo2 { 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]); + auto image_pe = ggml_ext_slice(ctx, 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 bundle = [&](int i) { return ggml_ext_slice(ctx, 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; @@ -326,7 +324,7 @@ namespace Fibo2 { 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 = ggml_ext_slice(ctx, 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); diff --git a/src/model/diffusion/flux.hpp b/src/model/diffusion/flux.hpp index f8f8305dd4..01b4609bc7 100644 --- a/src/model/diffusion/flux.hpp +++ b/src/model/diffusion/flux.hpp @@ -335,9 +335,9 @@ namespace Flux { x = mlp_0->forward(ctx, x); if (use_mlp_silu_act) { - x = ggml_ext_silu_act(ctx->ggml_ctx, x); + x = ggml_ext_silu_act(ctx, x); } else { - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); } x = mlp_2->forward(ctx, x); return x; @@ -686,11 +686,11 @@ namespace Flux { auto mlp = ggml_view_3d(ctx->ggml_ctx, qkv_mlp, mlp_hidden_dim * mlp_mult_factor, qkv_mlp->ne[1], qkv_mlp->ne[2], qkv_mlp->nb[1], qkv_mlp->nb[2], hidden_size * 3 * qkv_mlp->nb[0]); if (use_yak_mlp) { - mlp = ggml_ext_silu_act(ctx->ggml_ctx, mlp, false); + mlp = ggml_ext_silu_act(ctx, mlp, false); } else if (use_mlp_silu_act) { - mlp = ggml_ext_silu_act(ctx->ggml_ctx, mlp); + mlp = ggml_ext_silu_act(ctx, mlp); } else { - mlp = ggml_ext_gelu(ctx->ggml_ctx, mlp, true); + mlp = ggml_ext_gelu(ctx, mlp, true); } auto attn_mlp = ggml_concat(ctx->ggml_ctx, attn, mlp, 0); // [N, n_token, hidden_size + mlp_hidden_dim] auto output = linear2->forward(ctx, attn_mlp); // [N, n_token, hidden_size] @@ -745,7 +745,7 @@ namespace Flux { auto adaLN_modulation_1 = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"]); auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 2 * hidden_size] - auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0); + auto m_vec = ggml_ext_chunk(ctx, m, 2, 0); shift = m_vec[0]; scale = m_vec[1]; } @@ -832,16 +832,16 @@ namespace Flux { int64_t hidden_size_x = x->ne[0]; auto mlp_params = param_generator->forward(ctx, s); - auto fc_params = ggml_ext_chunk(ctx->ggml_ctx, mlp_params, 3, 0); + auto fc_params = ggml_ext_chunk(ctx, mlp_params, 3, 0); auto fc1_gate = ggml_reshape_3d(ctx->ggml_ctx, fc_params[0], hidden_size_x * mlp_ratio, hidden_size_x, batch_size); auto fc1_value = ggml_reshape_3d(ctx->ggml_ctx, fc_params[1], hidden_size_x * mlp_ratio, hidden_size_x, batch_size); auto fc2 = ggml_reshape_3d(ctx->ggml_ctx, fc_params[2], hidden_size_x, mlp_ratio * hidden_size_x, batch_size); - fc1_gate = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc1_gate, 1, 0, 2, 3)); // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x] + fc1_gate = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc1_gate, 1, 0, 2, 3)); // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x] fc1_gate = ggml_l2_norm(ctx->ggml_ctx, fc1_gate, 1e-12f); - fc1_value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc1_value, 1, 0, 2, 3)); // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x] + fc1_value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc1_value, 1, 0, 2, 3)); // [batch_size, hidden_size_x*mlp_ratio, hidden_size_x] fc1_value = ggml_l2_norm(ctx->ggml_ctx, fc1_value, 1e-12f); - fc2 = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, fc2, 1, 0, 2, 3)); // [batch_size, hidden_size_x, hidden_size_x*mlp_ratio] + fc2 = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, fc2, 1, 0, 2, 3)); // [batch_size, hidden_size_x, hidden_size_x*mlp_ratio] fc2 = ggml_l2_norm(ctx->ggml_ctx, fc2, 1e-12f); auto res_x = x; @@ -892,9 +892,9 @@ namespace Flux { auto norm = std::dynamic_pointer_cast(blocks["norm"]); auto conv = std::dynamic_pointer_cast(blocks["conv"]); - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [N, H, W, C] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [N, H, W, C] x = norm->forward(ctx, x); - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); // [N, C, H, W] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, H, W] x = conv->forward(ctx, x); return x; @@ -1028,13 +1028,13 @@ namespace Flux { if (config.is_chroma) { int64_t mod_index_length = 344; auto approx = std::dynamic_pointer_cast(blocks["distilled_guidance_layer"]); - auto distill_timestep = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 16, 10000, 1000.f); - auto distill_guidance = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 16, 10000, 1000.f); + auto distill_timestep = ggml_ext_timestep_embedding(ctx, timesteps, 16, 10000, 1000.f); + auto distill_guidance = ggml_ext_timestep_embedding(ctx, guidance, 16, 10000, 1000.f); // auto mod_index_arange = ggml_arange(ctx, 0, (float)mod_index_length, 1); // ggml_arange tot working on a lot of backends, precomputing it on CPU instead GGML_ASSERT(mod_index_arange != nullptr); - auto modulation_index = ggml_ext_timestep_embedding(ctx->ggml_ctx, mod_index_arange, 32, 10000, 1000.f); // [1, 344, 32] + auto modulation_index = ggml_ext_timestep_embedding(ctx, mod_index_arange, 32, 10000, 1000.f); // [1, 344, 32] // Batch broadcast (will it ever be useful) modulation_index = ggml_repeat(ctx->ggml_ctx, modulation_index, ggml_new_tensor_3d(ctx->ggml_ctx, GGML_TYPE_F32, modulation_index->ne[0], modulation_index->ne[1], img->ne[2])); // [N, 344, 32] @@ -1057,12 +1057,12 @@ namespace Flux { vec = dual_time_embed->forward(ctx, timestep_sem, timestep_tex); } else { auto time_in = std::dynamic_pointer_cast(blocks["time_in"]); - vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f)); + vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1000.f)); if (config.guidance_embed) { GGML_ASSERT(guidance != nullptr); auto guidance_in = std::dynamic_pointer_cast(blocks["guidance_in"]); // bf16 and fp16 result is different - auto g_in = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1000.f); + auto g_in = ggml_ext_timestep_embedding(ctx, guidance, 256, 10000, 1000.f); vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, g_in)); } @@ -1223,7 +1223,7 @@ namespace Flux { img = img_in_patch->forward(ctx, img); // [N, hidden_size, H/patch_size, W/patch_size] img = ggml_reshape_3d(ctx->ggml_ctx, img, img->ne[0] * img->ne[1], img->ne[2], img->ne[3]); // [N, hidden_size, H/patch_size*W/patch_size] - img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); // [N, H/patch_size*W/patch_size, hidden_size] + img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3)); // [N, H/patch_size*W/patch_size, hidden_size] auto out = forward_orig(ctx, img, context, timestep, y, guidance, pe, mod_index_arange, skip_layers, pulid_id, pulid_id_weight); // [N, n_img_token, hidden_size] @@ -1232,14 +1232,14 @@ namespace Flux { auto nerf_image_embedder = std::dynamic_pointer_cast(blocks["nerf_image_embedder"]); auto nerf_final_layer_conv = std::dynamic_pointer_cast(blocks["nerf_final_layer_conv"]); - auto nerf_pixels = DiT::patchify(ctx->ggml_ctx, orig_img, patch_size, patch_size); // [N, num_patches, C * patch_size * patch_size] + auto nerf_pixels = DiT::patchify(ctx, orig_img, patch_size, patch_size); // [N, num_patches, C * patch_size * patch_size] int64_t num_patches = nerf_pixels->ne[1]; nerf_pixels = ggml_reshape_3d(ctx->ggml_ctx, nerf_pixels, nerf_pixels->ne[0] / C, C, - nerf_pixels->ne[1] * nerf_pixels->ne[2]); // [N*num_patches, C, patch_size*patch_size] - nerf_pixels = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, nerf_pixels, 1, 0, 2, 3)); // [N*num_patches, patch_size*patch_size, C] + nerf_pixels->ne[1] * nerf_pixels->ne[2]); // [N*num_patches, C, patch_size*patch_size] + nerf_pixels = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, nerf_pixels, 1, 0, 2, 3)); // [N*num_patches, patch_size*patch_size, C] auto nerf_hidden = ggml_reshape_2d(ctx->ggml_ctx, out, out->ne[0], out->ne[1] * out->ne[2]); // [N*num_patches, hidden_size] auto img_dct = nerf_image_embedder->forward(ctx, nerf_pixels, dct); // [N*num_patches, patch_size*patch_size, nerf_hidden_size] @@ -1250,9 +1250,9 @@ namespace Flux { img_dct = block->forward(ctx, img_dct, nerf_hidden); } - img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img_dct, 1, 0, 2, 3)); // [N*num_patches, nerf_hidden_size, patch_size*patch_size] + img_dct = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img_dct, 1, 0, 2, 3)); // [N*num_patches, nerf_hidden_size, patch_size*patch_size] img_dct = ggml_reshape_3d(ctx->ggml_ctx, img_dct, img_dct->ne[0] * img_dct->ne[1], num_patches, img_dct->ne[2] / num_patches); // [N, num_patches, nerf_hidden_size*patch_size*patch_size] - img_dct = DiT::unpatchify(ctx->ggml_ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size); // [N, nerf_hidden_size, H, W] + img_dct = DiT::unpatchify(ctx, img_dct, (H + pad_h) / patch_size, (W + pad_w) / patch_size, patch_size, patch_size); // [N, nerf_hidden_size, H, W] out = nerf_final_layer_conv->forward(ctx, img_dct); // [N, C, H, W] @@ -1331,7 +1331,7 @@ namespace Flux { out = ggml_cont(ctx->ggml_ctx, out); } - out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size); // [N, C, H, W] + out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size); // [N, C, H, W] return out; } diff --git a/src/model/diffusion/flux3.hpp b/src/model/diffusion/flux3.hpp index 105b22565d..ad86d8ad69 100644 --- a/src/model/diffusion/flux3.hpp +++ b/src/model/diffusion/flux3.hpp @@ -95,7 +95,7 @@ namespace Flux3 { auto k = norm->key_norm(ctx, split(linear(ctx, "k_proj", x_mod))); auto v = split(linear(ctx, "v_proj", x_mod)); auto attn = Rope::attention(ctx, q, k, v, pe, nullptr); - auto mlp = ggml_ext_silu_act(gc, linear(ctx, "mlp_in", x_mod)); + auto mlp = ggml_ext_silu_act(ctx, linear(ctx, "mlp_in", x_mod)); auto update = ggml_add(gc, linear(ctx, "attn_out", attn), linear(ctx, "mlp_out", mlp)); ggml_tensor* out = nullptr; offset = 0; @@ -140,16 +140,16 @@ namespace Flux3 { auto gc = ctx->ggml_ctx; const int64_t width = image->ne[0], height = image->ne[1]; auto time_in = std::dynamic_pointer_cast(blocks["time_in"]); - auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, timestep, 256, 10000, 1000.f)); - auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(gc, zero, 256, 10000, 1000.f)); + auto image_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1000.f)); + auto clean_vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, zero, 256, 10000, 1000.f)); if (config.flow_map) { auto ref_time = std::dynamic_pointer_cast(blocks["ref_time_in"]); - auto ref_vec = ref_time->forward(ctx, ggml_ext_timestep_embedding(gc, target, 256, 10000, 1000.f)); + auto ref_vec = ref_time->forward(ctx, ggml_ext_timestep_embedding(ctx, target, 256, 10000, 1000.f)); image_vec = ggml_add(gc, image_vec, ref_vec); clean_vec = ggml_add(gc, clean_vec, ref_vec); } std::vector modes{"image"}; - std::vector inputs{DiT::patchify(gc, image, 1, 1)}; + std::vector inputs{DiT::patchify(ctx, image, 1, 1)}; if (reference) { modes.push_back("image_cond"); inputs.push_back(reference); @@ -184,7 +184,7 @@ namespace Flux3 { 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); + return DiT::unpatchify(ctx, out, height, width, 1, 1); } }; @@ -215,9 +215,10 @@ namespace Flux3 { GGML_ASSERT(params.x && params.context && params.timesteps); GGML_ASSERT(!config.flow_map || extra->target_timestep); auto get_graph = [&]() { - auto gf = new_graph_custom(FLUX_GRAPH_SIZE); - auto x = make_input(*params.x); - auto text = make_input(*params.context); + auto gf = new_graph_custom(FLUX_GRAPH_SIZE); + auto runner_ctx = get_context(); + auto x = make_input(*params.x); + auto text = make_input(*params.context); GGML_ASSERT(x->ne[3] == 1 && text->ne[2] == 1); std::vector refs; ggml_tensor* references = nullptr; @@ -225,13 +226,13 @@ namespace Flux3 { for (const auto& ref : *params.ref_latents) { auto r = make_input(ref); refs.push_back(r); - r = DiT::patchify(compute_ctx, r, 1, 1); + r = DiT::patchify(&runner_ctx, r, 1, 1); references = references ? ggml_concat(compute_ctx, references, r, 1) : r; } } 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, [&]() { + 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); @@ -254,9 +255,8 @@ namespace Flux3 { 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); + 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); 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 90bfb51b34..8581fba571 100644 --- a/src/model/diffusion/hidream_o1.hpp +++ b/src/model/diffusion/hidream_o1.hpp @@ -197,7 +197,7 @@ namespace HiDreamO1 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) { auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); - auto emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1000.0f); + auto emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10000, 1000.0f); emb = mlp_0->forward(ctx, emb); emb = ggml_silu_inplace(ctx->ggml_ctx, emb); emb = mlp_2->forward(ctx, emb); @@ -292,7 +292,6 @@ namespace HiDreamO1 { ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { return LLM::LLMRunner::encode_image_common(this, - compute_ctx, runner_ctx, image, config.llm.vision, @@ -398,7 +397,7 @@ namespace HiDreamO1 { auto t_emb = t_embedder1->forward(&runner_ctx, timestep); 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); + auto prefix = ggml_ext_slice(&runner_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); } else { txt = ggml_reshape_3d(compute_ctx, t_emb, t_emb->ne[0], 1, 1); @@ -425,8 +424,8 @@ namespace HiDreamO1 { } x_pred_start = first_vinput; } - auto x_pred = ggml_ext_slice(compute_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens); - x_pred = DiT::unpatchify_and_crop(compute_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE); + auto x_pred = ggml_ext_slice(&runner_ctx, x_pred_all, 1, x_pred_start, x_pred_start + target_tokens); + x_pred = DiT::unpatchify_and_crop(&runner_ctx, x_pred, x->ne[1], x->ne[0], PATCH_SIZE, PATCH_SIZE); float sigma = 1.0f - timestep_tensor.values()[0]; sigma = std::max(1e-6f, sigma); diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index 5fc29b6b0a..d2bc8b5804 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -50,7 +50,7 @@ namespace Hunyuan { // self attn auto qkv = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt)); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); auto q = qkv_vec[0]; auto k = qkv_vec[1]; auto v = qkv_vec[2]; @@ -60,7 +60,7 @@ namespace Hunyuan { // adaLN_modulation auto emb = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, t_emb)); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0); + auto mods = ggml_ext_chunk(ctx, emb, 2, 0); txt = ggml_add(ctx->ggml_ctx, txt, ggml_mul(ctx->ggml_ctx, attn_out, mods[0])); @@ -122,11 +122,11 @@ namespace Hunyuan { auto c_embedder = std::dynamic_pointer_cast(blocks["c_embedder"]); auto individual_token_refiner = std::dynamic_pointer_cast(blocks["individual_token_refiner"]); - auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f)); + auto t_emb = t_embedder->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f)); - auto h = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, txt, 1, 0, 2, 3)); + auto h = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, txt, 1, 0, 2, 3)); auto pooled_projections = ggml_scale(ctx->ggml_ctx, ggml_sum_rows(ctx->ggml_ctx, h), 1.f / txt->ne[1]); - pooled_projections = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, pooled_projections, 1, 0, 2, 3)); + pooled_projections = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, pooled_projections, 1, 0, 2, 3)); auto c_emb = c_embedder->forward(ctx, pooled_projections); t_emb = ggml_add(ctx->ggml_ctx, t_emb, c_emb); @@ -151,9 +151,9 @@ namespace Hunyuan { auto fc3 = std::dynamic_pointer_cast(blocks["fc3"]); x = fc1->forward(ctx, layernorm->forward(ctx, x)); - x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = ggml_ext_gelu(ctx, x); x = fc2->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = ggml_ext_gelu(ctx, x); return fc3->forward(ctx, x); } }; @@ -364,11 +364,12 @@ namespace Hunyuan { return x; } - ggml_tensor* unpatchify(struct ggml_context* ctx, + ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t t_len, int64_t h_len, int64_t w_len) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, t_len*h_len*w_len, C*pt*ph*pw] // return: [N*C, t_len*pt, h_len*ph, w_len*pw] int64_t N = x->ne[3]; @@ -379,15 +380,15 @@ namespace Hunyuan { GGML_ASSERT(C * pt * ph * pw == x->ne[0]); - x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] - x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] return x; } @@ -431,10 +432,10 @@ namespace Hunyuan { img = img_in->forward(ctx, img); // [N*C, t_len*h_len*w_len, hidden_size] txt = txt_in->forward(ctx, txt, timestep, nullptr); // [N, n_txt_token, hidden_size] - auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f)); + auto vec = time_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f)); if (config.use_meanflow && timestep_r != nullptr) { auto time_r_in = std::dynamic_pointer_cast(blocks["time_r_in"]); - auto vec_r = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_r, 256, 10000, 1000.f)); + auto vec_r = time_r_in->forward(ctx, ggml_ext_timestep_embedding(ctx, timestep_r, 256, 10000, 1000.f)); vec = ggml_add(ctx->ggml_ctx, vec, vec_r); if (!config.use_meanflow_sum) { vec = ggml_scale(ctx->ggml_ctx, vec, 0.5f); @@ -446,7 +447,7 @@ namespace Hunyuan { } if (config.guidance_embed && guidance != nullptr) { auto guidance_in = std::dynamic_pointer_cast(blocks["guidance_in"]); - auto guidance_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, guidance, 256, 10000, 1.f); + auto guidance_emb = ggml_ext_timestep_embedding(ctx, guidance, 256, 10000, 1.f); vec = ggml_add(ctx->ggml_ctx, vec, guidance_in->forward(ctx, guidance_emb)); } @@ -529,12 +530,12 @@ namespace Hunyuan { auto out = forward_orig(ctx, x, context, timestep, pe, guidance, y, txt_byt5, clip_fea, timestep_r, N); - out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] + out = unpatchify(ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] // slice - out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] + out = ggml_ext_slice(ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] + out = ggml_ext_slice(ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] + out = ggml_ext_slice(ctx, out, 0, 0, W); // [N*C, T, H, W] return out; } diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index 270344dc97..fbc189aa3b 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -73,12 +73,13 @@ namespace Ideogram4 { } }; - __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* timestep_embedding_sin_cos(GGMLRunnerContext* runner_ctx, ggml_tensor* timesteps, int dim) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(dim % 2 == 0); - auto embedding = ggml_ext_timestep_embedding(ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f); - auto chunks = ggml_ext_chunk(ctx, embedding, 2, 0); + auto embedding = ggml_ext_timestep_embedding(runner_ctx, timesteps, dim, TIMESTEP_MAX_PERIOD, 10.f); + auto chunks = ggml_ext_chunk(runner_ctx, embedding, 2, 0); return ggml_concat(ctx, chunks[1], chunks[0], 0); } @@ -108,7 +109,10 @@ namespace Ideogram4 { return ggml_ext_modulate(ctx, x, nullptr, scale); } - __STATIC_INLINE__ ggml_tensor* patchify(ggml_context* ctx, ggml_tensor* x, const Ideogram4Config& config) { + __STATIC_INLINE__ ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + const Ideogram4Config& config) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, 128, H, W] with channel order [ae, ph, pw]. // return: [N, H*W, 128] with token channel order [ph, pw, ae]. const int64_t W = x->ne[0]; @@ -126,11 +130,12 @@ namespace Ideogram4 { return x; } - __STATIC_INLINE__ ggml_tensor* unpatchify(ggml_context* ctx, + __STATIC_INLINE__ ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t H, int64_t W, const Ideogram4Config& config) { + auto ctx = runner_ctx->ggml_ctx; const int64_t C = x->ne[0]; const int64_t N = x->ne[2]; @@ -217,7 +222,7 @@ namespace Ideogram4 { auto out_proj = std::dynamic_pointer_cast(blocks["o"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N); auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N); auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N); @@ -278,7 +283,7 @@ namespace Ideogram4 { auto adaln_modulation = std::dynamic_pointer_cast(blocks["adaln_modulation"]); auto mod = adaln_modulation->forward(ctx, adaln_input); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 4, 0); + auto mods = ggml_ext_chunk(ctx, mod, 4, 0); auto scale_msa = mods[0]; auto gate_msa = to_token_modulation(ctx->ggml_ctx, ggml_tanh(ctx->ggml_ctx, mods[1])); auto scale_mlp = mods[2]; @@ -315,7 +320,7 @@ namespace Ideogram4 { auto mlp_in = std::dynamic_pointer_cast(blocks["mlp_in"]); auto mlp_out = std::dynamic_pointer_cast(blocks["mlp_out"]); - x = timestep_embedding_sin_cos(ctx->ggml_ctx, x, static_cast(dim)); + x = timestep_embedding_sin_cos(ctx, x, static_cast(dim)); x = ggml_silu(ctx->ggml_ctx, mlp_in->forward(ctx, x)); x = mlp_out->forward(ctx, x); return x; @@ -383,7 +388,7 @@ namespace Ideogram4 { auto embed_image_indicator = std::dynamic_pointer_cast(blocks["embed_image_indicator"]); auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); - auto img = patchify(ctx->ggml_ctx, x, config); + auto img = patchify(ctx, x, config); img = input_proj->forward(ctx, img); ggml_tensor* h = img; @@ -413,11 +418,11 @@ namespace Ideogram4 { h = final_layer->forward(ctx, h, adaln_input); if (context_len > 0) { - h = ggml_ext_slice(ctx->ggml_ctx, h, 1, context_len, h->ne[1]); + h = ggml_ext_slice(ctx, h, 1, context_len, h->ne[1]); } - h = unpatchify(ctx->ggml_ctx, h, H, W, config); - h = ggml_ext_scale(ctx->ggml_ctx, h, -1.f); + h = unpatchify(ctx, h, H, W, config); + h = ggml_ext_scale(ctx, h, -1.f); return h; } }; diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp index 3681c65ab2..92f4b2e207 100644 --- a/src/model/diffusion/iris.hpp +++ b/src/model/diffusion/iris.hpp @@ -92,7 +92,7 @@ namespace Iris { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) { auto c = ctx->ggml_ctx; - auto qkv = split_qkv(c, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x)); + auto qkv = ggml_ext_split_qkv(ctx, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x)); auto q = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]); auto k = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]); auto v = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]); @@ -247,24 +247,24 @@ namespace Iris { std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* y, ggml_tensor* img_mod, ggml_tensor* txt_mod, ggml_tensor* pe) { auto c = ctx->ggml_ctx; int64_t nt = y->ne[1]; - auto mx = ggml_ext_chunk(c, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0); + auto mx = ggml_ext_chunk(ctx, ggml_add(c, img_mod, params[dual ? "adaln_img.bias" : "adaln.bias"]), 6, 0); if (!dual) { auto tokens = ggml_concat(c, y, x, 1); auto h = Pid::apply_adaln(c, norm(ctx, "norm1", tokens), mx[0], mx[1]); auto qkv = project(ctx, h, ""); auto out = Rope::attention(ctx, qkv[0], qkv[1], qkv[2], pe, nullptr); tokens = finish(ctx, tokens, h, out, mx, ""); - return {ggml_ext_slice(c, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(c, tokens, 1, 0, nt)}; + return {ggml_ext_slice(ctx, tokens, 1, nt, tokens->ne[1]), ggml_ext_slice(ctx, tokens, 1, 0, nt)}; } - auto my = ggml_ext_chunk(c, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0); + auto my = ggml_ext_chunk(ctx, ggml_add(c, txt_mod, params["adaln_txt.bias"]), 6, 0); auto hx = Pid::apply_adaln(c, norm(ctx, "norm_x1", x), mx[0], mx[1]); auto hy = Pid::apply_adaln(c, norm(ctx, "norm_y1", y), my[0], my[1]); auto qx = project(ctx, hx, "_x"); auto qy = project(ctx, hy, "_y"); auto out = Rope::attention(ctx, ggml_concat(c, qy[0], qx[0], 2), ggml_concat(c, qy[1], qx[1], 2), ggml_concat(c, qy[2], qx[2], 2), pe, nullptr); - x = finish(ctx, x, hx, ggml_ext_slice(c, out, 1, nt, out->ne[1]), mx, "_x"); - y = finish(ctx, y, hy, ggml_ext_slice(c, out, 1, 0, nt), my, "_y"); + x = finish(ctx, x, hx, ggml_ext_slice(ctx, out, 1, nt, out->ne[1]), mx, "_x"); + y = finish(ctx, y, hy, ggml_ext_slice(ctx, out, 1, 0, nt), my, "_y"); return {x, y}; } }; @@ -294,7 +294,7 @@ namespace Iris { auto batch = cond->ne[2]; auto mods = std::dynamic_pointer_cast(blocks["adaln"])->forward(ctx, cond); mods = ggml_reshape_3d(c, mods, 4 * d, p2, patches * batch); - auto m = ggml_ext_chunk(c, mods, 4, 0); + auto m = ggml_ext_chunk(ctx, mods, 4, 0); auto h = std::dynamic_pointer_cast(blocks["norm1"])->forward(ctx, x); h = ggml_reshape_3d(c, h, d * p2, patches, batch); h = std::dynamic_pointer_cast(blocks["compress"])->forward(ctx, h); @@ -337,7 +337,7 @@ namespace Iris { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* image, ggml_tensor* t, ggml_tensor* text, ggml_tensor* mask, ggml_tensor* joint_pe, ggml_tensor* pixel_pe, ggml_tensor* pixel_pos) { auto c = ctx->ggml_ctx; int p = static_cast(config.patch_size); - auto x = DiT::patchify(c, image, p, p, true); + auto x = DiT::patchify(ctx, image, p, p, true); x = std::dynamic_pointer_cast(blocks["s_embedder"])->forward(ctx, x); auto te = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); te = ggml_reshape_3d(c, te, config.hidden_size, 1, image->ne[3]); @@ -360,7 +360,7 @@ namespace Iris { } x = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, x); x = ggml_reshape_3d(c, x, 3 * p * p, cond->ne[1], image->ne[3]); - return DiT::unpatchify(c, x, image->ne[1] / p, image->ne[0] / p, p, p, false); + return DiT::unpatchify(ctx, x, image->ne[1] / p, image->ne[0] / p, p, p, false); } }; diff --git a/src/model/diffusion/krea2.hpp b/src/model/diffusion/krea2.hpp index e44c5a9c6e..54dd6fdc44 100644 --- a/src/model/diffusion/krea2.hpp +++ b/src/model/diffusion/krea2.hpp @@ -185,9 +185,9 @@ namespace Krea2 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { ggml_tensor* scale = params["scale"]; if (ctx->weight_adapter) { - scale = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, scale, prefix + "scale.weight"); + scale = ctx->weight_adapter->patch_weight(ctx, scale, prefix + "scale.weight"); } - scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx->ggml_ctx, scale->ne[0], 1, 1, 1)); + scale = ggml_add(ctx->ggml_ctx, scale, ggml_ext_ones(ctx, scale->ne[0], 1, 1, 1)); x = ggml_rms_norm(ctx->ggml_ctx, x, eps); x = ggml_mul_inplace(ctx->ggml_ctx, x, scale); return x; @@ -273,9 +273,7 @@ namespace Krea2 { auto knorm = std::dynamic_pointer_cast(blocks["qknorm.knorm"]); auto wo = std::dynamic_pointer_cast(blocks["wo"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - wo->set_force_prec_f32(true); - } + wo->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); int64_t L = x->ne[1]; int64_t N = x->ne[2]; @@ -316,11 +314,11 @@ namespace Krea2 { std::vector forward(GGMLRunnerContext* ctx, ggml_tensor* vec) { auto lin = params["lin"]; if (ctx->weight_adapter) { - lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight"); + lin = ctx->weight_adapter->patch_weight(ctx, lin, prefix + "lin.weight"); } lin = ggml_repeat(ctx->ggml_ctx, lin, vec); auto out = ggml_add(ctx->ggml_ctx, vec, lin); - return ggml_ext_chunk(ctx->ggml_ctx, out, 6, 0); + return ggml_ext_chunk(ctx, out, 6, 0); } }; @@ -342,10 +340,10 @@ namespace Krea2 { std::vector forward(GGMLRunnerContext* ctx, ggml_tensor* vec) { auto lin = params["lin"]; if (ctx->weight_adapter) { - lin = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, lin, prefix + "lin.weight"); + lin = ctx->weight_adapter->patch_weight(ctx, lin, prefix + "lin.weight"); } auto out = ggml_add(ctx->ggml_ctx, lin, vec); - return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1); + return ggml_ext_chunk(ctx, out, 2, 1); } }; @@ -541,7 +539,7 @@ namespace Krea2 { auto linear_0 = std::dynamic_pointer_cast(blocks["0"]); auto linear_2 = std::dynamic_pointer_cast(blocks["2"]); x = linear_0->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, false); + x = ggml_ext_gelu(ctx, x, false); x = linear_2->forward(ctx, x); return x; } @@ -555,7 +553,7 @@ namespace Krea2 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { auto linear_1 = std::dynamic_pointer_cast(blocks["1"]); - x = ggml_ext_gelu(ctx->ggml_ctx, x, false); + x = ggml_ext_gelu(ctx, x, false); x = linear_1->forward(ctx, x); return x; } @@ -575,7 +573,7 @@ namespace Krea2 { auto linear_3 = std::dynamic_pointer_cast(blocks["3"]); x = norm->forward(ctx, x); x = linear_1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = linear_3->forward(ctx, x); return x; } @@ -656,7 +654,7 @@ namespace Krea2 { int64_t ref_len = img->ne[1] - img_len; img = first->forward(ctx, img); - auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, static_cast(config.timestep_dim), 10000, 1000.f); + auto t = ggml_ext_timestep_embedding(ctx, timestep, static_cast(config.timestep_dim), 10000, 1000.f); t = tmlp->forward(ctx, t); t = ggml_reshape_3d(ctx->ggml_ctx, t, t->ne[0], 1, t->ne[1]); auto tvec = tproj->forward(ctx, t); @@ -665,7 +663,7 @@ namespace Krea2 { if (ref_latents.size() > 0 && zero_timestep_refs) { // "index_timestep_zero" mode: use timestep = 0 for ref latents auto timestep_0 = ggml_scale(ctx->ggml_ctx, timestep, 0.0f); - auto t_0 = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_0, static_cast(config.timestep_dim), 10000, 1000.f); + auto t_0 = ggml_ext_timestep_embedding(ctx, timestep_0, static_cast(config.timestep_dim), 10000, 1000.f); t_0 = tmlp->forward(ctx, t_0); t_0 = ggml_reshape_3d(ctx->ggml_ctx, t_0, t_0->ne[0], 1, t_0->ne[1]); tvec_0 = tproj->forward(ctx, t_0); @@ -683,9 +681,9 @@ namespace Krea2 { sd::ggml_graph_cut::mark_graph_cut(hidden_states, "krea2.blocks." + std::to_string(i), "hidden_states"); } - hidden_states = ggml_ext_slice(ctx->ggml_ctx, hidden_states, 1, txt_len, txt_len + img_len); + hidden_states = ggml_ext_slice(ctx, hidden_states, 1, txt_len, txt_len + img_len); hidden_states = last->forward(ctx, hidden_states, t); - hidden_states = DiT::unpatchify_and_crop(ctx->ggml_ctx, hidden_states, H, W, config.patch_size, config.patch_size, true); + hidden_states = DiT::unpatchify_and_crop(ctx, hidden_states, H, W, config.patch_size, config.patch_size, true); return hidden_states; } }; diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 0f8c816739..5c86e65d01 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -87,7 +87,7 @@ namespace Lens { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) { auto timestep_embedder = std::dynamic_pointer_cast(blocks["timestep_embedder"]); - auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1000.f); + auto timesteps_proj = ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1000.f); return timestep_embedder->forward(ctx, timesteps_proj); } }; @@ -151,8 +151,8 @@ namespace Lens { int64_t N = img->ne[2]; int64_t inner = dim_head * num_heads; - auto img_qkv_vec = split_qkv(ctx->ggml_ctx, img_qkv->forward(ctx, img)); - auto txt_qkv_vec = split_qkv(ctx->ggml_ctx, txt_qkv->forward(ctx, txt)); + auto img_qkv_vec = ggml_ext_split_qkv(ctx, img_qkv->forward(ctx, img)); + auto txt_qkv_vec = ggml_ext_split_qkv(ctx, txt_qkv->forward(ctx, txt)); auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N); auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N); @@ -232,9 +232,9 @@ namespace Lens { auto temb = ggml_silu(ctx->ggml_ctx, t_emb); auto img_mod_params = img_mod_1->forward(ctx, temb); - auto img_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, img_mod_params, 6, 0); + auto img_mod_vec = ggml_ext_chunk(ctx, img_mod_params, 6, 0); auto txt_mod_params = txt_mod_1->forward(ctx, temb); - auto txt_mod_vec = ggml_ext_chunk(ctx->ggml_ctx, txt_mod_params, 6, 0); + auto txt_mod_vec = ggml_ext_chunk(ctx, txt_mod_params, 6, 0); auto img_normed = img_norm1->forward(ctx, img); auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_vec[0], img_mod_vec[1]); @@ -268,7 +268,7 @@ namespace Lens { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* conditioning) { auto linear = std::dynamic_pointer_cast(blocks["linear"]); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0); + auto mods = ggml_ext_chunk(ctx, linear->forward(ctx, ggml_silu(ctx->ggml_ctx, conditioning)), 2, 0); auto scale = mods[0]; auto shift = mods[1]; x = ggml_norm(ctx->ggml_ctx, x, eps); @@ -319,10 +319,10 @@ namespace Lens { auto t_emb = time_text_embed->forward(ctx, timestep); auto img = ggml_reshape_3d(ctx->ggml_ctx, x, W * H, C, N); - img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); + img = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3)); img = img_in->forward(ctx, img); - std::vector txt_chunks = ggml_ext_chunk(ctx->ggml_ctx, context, config.selected_layer_count, 0); + std::vector txt_chunks = ggml_ext_chunk(ctx, context, config.selected_layer_count, 0); ggml_tensor* txt = nullptr; for (int i = 0; i < config.selected_layer_count; ++i) { auto txt_norm = std::dynamic_pointer_cast(blocks["txt_norm." + std::to_string(i)]); @@ -346,7 +346,7 @@ namespace Lens { img = norm_out->forward(ctx, img, t_emb); img = proj_out->forward(ctx, img); - auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, img, 1, 0, 2, 3)); + auto out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, img, 1, 0, 2, 3)); out = ggml_reshape_4d(ctx->ggml_ctx, out, W, H, config.patch_size * config.patch_size * config.out_channels, N); return out; } diff --git a/src/model/diffusion/lingbot_video.hpp b/src/model/diffusion/lingbot_video.hpp index b94d7fdb86..fa6024edf3 100644 --- a/src/model/diffusion/lingbot_video.hpp +++ b/src/model/diffusion/lingbot_video.hpp @@ -365,7 +365,7 @@ namespace LingBotVideo { } selected_mask = ggml_reshape_2d(gctx, selected_mask, num_experts, n_token_total); - ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(gctx, selected_mask)); + ggml_tensor* excluded_group_mask = ggml_sub(gctx, selected_mask, ggml_ext_ones_like(ctx, selected_mask)); ggml_tensor* masked_scores = ggml_add(gctx, choice_scores, ggml_scale(gctx, excluded_group_mask, 1.0e9f)); return ggml_argsort_top_k(gctx, masked_scores, static_cast(num_experts_per_tok)); } @@ -476,7 +476,7 @@ namespace LingBotVideo { auto table = ggml_reshape_3d(ctx->ggml_ctx, params["scale_shift_table"], hidden_size * 6, 1, 1); auto mod = ggml_add(ctx->ggml_ctx, temb6, table); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, mod, 6, 0); + auto mods = ggml_ext_chunk(ctx, mod, 6, 0); auto shift_msa = mods[0]; auto scale_msa = mods[1]; @@ -556,13 +556,13 @@ namespace LingBotVideo { int64_t w_len = W / config.patch_w; int64_t n_img = t_len * h_len * w_len; - auto img = DiT::patchify_3d(ctx->ggml_ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false); + auto img = DiT::patchify_3d(ctx, x, config.patch_t, config.patch_h, config.patch_w, N, false); img = patch_embedder->forward(ctx, img); auto txt = text_embedder->forward(ctx, context); auto hidden_states = ggml_concat(ctx->ggml_ctx, img, txt, 1); - auto timestep_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, + auto timestep_proj = ggml_ext_timestep_embedding(ctx, timestep, static_cast(config.freq_dim), 10000, @@ -582,16 +582,16 @@ namespace LingBotVideo { sd::ggml_graph_cut::mark_graph_cut(hidden_states, "lingbot_video.blocks." + std::to_string(i), "hidden_states"); } - auto final_mods = ggml_ext_chunk(ctx->ggml_ctx, + auto final_mods = ggml_ext_chunk(ctx, norm_out_modulation->forward(ctx, ggml_silu(ctx->ggml_ctx, temb_tokens)), 2, 0); hidden_states = norm_out->forward(ctx, hidden_states); 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); + hidden_states = ggml_ext_slice(ctx, hidden_states, 1, 0, n_img); - auto out = DiT::unpatchify_3d(ctx->ggml_ctx, + auto out = DiT::unpatchify_3d(ctx, hidden_states, t_len, h_len, diff --git a/src/model/diffusion/llada_image.hpp b/src/model/diffusion/llada_image.hpp index 6b2b534732..ca127b8309 100644 --- a/src/model/diffusion/llada_image.hpp +++ b/src/model/diffusion/llada_image.hpp @@ -217,8 +217,8 @@ namespace LLaDAImage { GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]); - 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]); + auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]); + auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]); for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)]); @@ -246,7 +246,7 @@ namespace LLaDAImage { txt_img = final_layer->forward(ctx, txt_img, t_emb); - return ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); + return ggml_ext_slice(ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); } ggml_tensor* pad_stream(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pad_token) { @@ -308,8 +308,8 @@ namespace LLaDAImage { GGML_ASSERT(cap_len * 2 + img_len * 2 + sig_len == pe->ne[3]); - auto cap_pe = ggml_ext_slice(gctx, pe, 3, 0, cap_len * 2); - auto img_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2); + auto cap_pe = ggml_ext_slice(ctx, pe, 3, 0, cap_len * 2); + auto img_pe = ggml_ext_slice(ctx, pe, 3, cap_len * 2, cap_len * 2 + img_len * 2); auto img_adaln = ggml_concat(gctx, per_token(t_clean, img_len), per_token(t_noisy, img_len), 1); @@ -322,7 +322,7 @@ namespace LLaDAImage { img = block->forward(ctx, img, img_pe, nullptr, img_adaln); } if (sig != nullptr) { - auto sig_pe = ggml_ext_slice(gctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]); + auto sig_pe = ggml_ext_slice(ctx, pe, 3, cap_len * 2 + img_len * 2, pe->ne[3]); for (int i = 0; i < config.num_refiner_layers; i++) { auto block = std::dynamic_pointer_cast(blocks["sigvq_refiner." + std::to_string(i)]); sig = block->forward(ctx, sig, sig_pe, nullptr, nullptr); @@ -350,7 +350,7 @@ namespace LLaDAImage { // The stream is padded to SEQ_MULTI_OF, so drop the pad tokens: they are not part of // the latent grid that unpatchify reconstructs. int64_t target_start = cap_len * 2 + img_len; - return ggml_ext_slice(gctx, seq, 1, target_start, target_start + n_img_token); + return ggml_ext_slice(ctx, seq, 1, target_start, target_start + n_img_token); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -372,10 +372,10 @@ namespace LLaDAImage { auto out = forward_core(ctx, img, timestep, context, pe); - out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, patch_size, patch_size, false); + out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size, false); // The reference pipeline negates the model output before the scheduler step. - return ggml_ext_scale(ctx->ggml_ctx, out, -1.f); + return ggml_ext_scale(ctx, out, -1.f); } }; @@ -488,8 +488,8 @@ namespace LLaDAImage { auto src = DiT::pad_and_patchify(&runner_ctx, source, config.patch_size, config.patch_size, false); auto out = llada_image.forward_editing(&runner_ctx, target, timesteps, context, semantic, src, pe); - out = DiT::unpatchify_and_crop(runner_ctx.ggml_ctx, out, H, W, config.patch_size, config.patch_size, false); - out = ggml_ext_scale(runner_ctx.ggml_ctx, out, -1.f); + out = DiT::unpatchify_and_crop(&runner_ctx, out, H, W, config.patch_size, config.patch_size, false); + out = ggml_ext_scale(&runner_ctx, out, -1.f); ggml_build_forward_expand(gf, out); return gf; diff --git a/src/model/diffusion/ltxv.hpp b/src/model/diffusion/ltxv.hpp index f5d2b36157..8a07e79ce8 100644 --- a/src/model/diffusion/ltxv.hpp +++ b/src/model/diffusion/ltxv.hpp @@ -581,7 +581,7 @@ namespace LTXV { auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); - auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, frequency_embedding_size); + auto t_emb = ggml_ext_timestep_embedding(ctx, timestep, frequency_embedding_size); t_emb = linear_1->forward(ctx, t_emb); t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb); t_emb = linear_2->forward(ctx, t_emb); @@ -631,7 +631,7 @@ namespace LTXV { auto linear_2 = std::dynamic_pointer_cast(blocks["linear_2"]); caption = linear_1->forward(ctx, caption); - caption = ggml_ext_gelu(ctx->ggml_ctx, caption, true); + caption = ggml_ext_gelu(ctx, caption, true); caption = linear_2->forward(ctx, caption); return caption; } @@ -650,7 +650,7 @@ namespace LTXV { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* caption) { auto linear_1 = std::dynamic_pointer_cast(blocks["linear_1"]); caption = ggml_rms_norm(ctx->ggml_ctx, caption, 1e-6f); - caption = ggml_ext_scale(ctx->ggml_ctx, caption, std::sqrt(static_cast(hidden_size) / static_cast(in_features))); + caption = ggml_ext_scale(ctx, caption, std::sqrt(static_cast(hidden_size) / static_cast(in_features))); return linear_1->forward(ctx, caption); } }; @@ -724,7 +724,7 @@ namespace LTXV { auto to_gate_logits = std::dynamic_pointer_cast(blocks["to_gate_logits"]); auto gate_logits = to_gate_logits->forward(ctx, x); auto gates = ggml_sigmoid(ctx->ggml_ctx, gate_logits); - gates = ggml_ext_scale(ctx->ggml_ctx, gates, 2.0f, true); + gates = ggml_ext_scale(ctx, gates, 2.0f, true); gates = ggml_reshape_4d(ctx->ggml_ctx, gates, 1, heads, gate_logits->ne[1], gate_logits->ne[2]); auto out4 = ggml_reshape_4d(ctx->ggml_ctx, out, dim_head, heads, out->ne[1], out->ne[2]); @@ -783,7 +783,7 @@ namespace LTXV { s = ggml_repeat(ctx->ggml_ctx, s, e); t = ggml_repeat(ctx->ggml_ctx, t, e); auto out = ggml_add(ctx->ggml_ctx, s, t); - return ggml_ext_chunk(ctx->ggml_ctx, out, static_cast(coeff), 1); + return ggml_ext_chunk(ctx, out, static_cast(coeff), 1); } std::vector get_prompt_scale_shift_values(GGMLRunnerContext* ctx, @@ -797,7 +797,7 @@ namespace LTXV { s = ggml_repeat(ctx->ggml_ctx, s, e); t = ggml_repeat(ctx->ggml_ctx, t, e); auto out = ggml_add(ctx->ggml_ctx, s, t); - return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1); + return ggml_ext_chunk(ctx, out, 2, 1); } ggml_tensor* forward(GGMLRunnerContext* ctx, @@ -945,7 +945,7 @@ namespace LTXV { for (int64_t i = 1; i < duplications; i++) { regs_full = ggml_concat(ctx->ggml_ctx, regs_full, regs, 1); } - regs_full = ggml_ext_slice(ctx->ggml_ctx, regs_full, 1, seq_len, seq_len + total_to_keep); + regs_full = ggml_ext_slice(ctx, regs_full, 1, seq_len, seq_len + total_to_keep); return ggml_concat(ctx->ggml_ctx, hidden_states, regs_full, 1); } @@ -1182,7 +1182,7 @@ namespace LTXV { t = ggml_repeat(ctx->ggml_ctx, t, e); s = ggml_repeat(ctx->ggml_ctx, s, e); auto out = ggml_add(ctx->ggml_ctx, s, t); - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, out, static_cast(coeff), 1); + auto chunks = ggml_ext_chunk(ctx, out, static_cast(coeff), 1); return std::vector(chunks.begin() + start, chunks.begin() + start + count); } @@ -1287,27 +1287,27 @@ namespace LTXV { auto ax_norm3 = rms_norm(ctx->ggml_ctx, ax); if (run_a2v) { - auto a2v_audio_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4); - auto a2v_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4); + auto a2v_audio_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4); + auto a2v_video_table = ggml_ext_slice(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, 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_table = ggml_ext_slice(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]; vx = ggml_add(ctx->ggml_ctx, vx, apply_gate(ctx->ggml_ctx, a2v_out, a2v_gate)); } if (run_v2a) { - auto v2a_audio_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4); - auto v2a_video_table = ggml_ext_slice(ctx->ggml_ctx, params["scale_shift_table_a2v_ca_video"], 1, 0, 4); + auto v2a_audio_table = ggml_ext_slice(ctx, params["scale_shift_table_a2v_ca_audio"], 1, 0, 4); + auto v2a_video_table = ggml_ext_slice(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, 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_table = ggml_ext_slice(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]; ax = ggml_add(ctx->ggml_ctx, ax, apply_gate(ctx->ggml_ctx, v2a_out, v2a_gate)); } @@ -1427,7 +1427,7 @@ namespace LTXV { ggml_tensor* patchify_video(GGMLRunnerContext* ctx, ggml_tensor* x, int64_t n) { x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / n, n); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); return x; } @@ -1436,14 +1436,14 @@ namespace LTXV { int64_t width, int64_t height, int64_t frames) { - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); x = ggml_reshape_4d(ctx->ggml_ctx, x, width, height, frames, x->ne[1] * x->ne[2]); return x; } ggml_tensor* patchify_audio(GGMLRunnerContext* ctx, ggml_tensor* ax) { // ax: [b, c, t, f] - ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3)); // [b, t, c, f] + ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, ax, 0, 2, 1, 3)); // [b, t, c, f] ax = ggml_reshape_3d(ctx->ggml_ctx, ax, ax->ne[0] * ax->ne[1], ax->ne[2], ax->ne[3]); // [b, t, c*f] return ax; } @@ -1462,7 +1462,7 @@ namespace LTXV { return nullptr; } ax = ggml_reshape_4d(ctx->ggml_ctx, ax, config.audio_frequency_bins, config.num_audio_channels, audio_length, ax->ne[2]); // [b, t, c, f] - ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ax, 0, 2, 1, 3)); // [b, c, t, f] + ax = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, ax, 0, 2, 1, 3)); // [b, c, t, f] return ax; } @@ -1483,10 +1483,10 @@ namespace LTXV { context->ne[1] < 1024; if (is_fully_processed_context) { - auto v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim); + auto v_context = ggml_ext_slice(ctx, context, 0, 0, config.cross_attention_dim); ggml_tensor* a_context = nullptr; if (process_audio_context) { - a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim); + a_context = ggml_ext_slice(ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim); } return {v_context, a_context}; } @@ -1494,14 +1494,14 @@ namespace LTXV { ggml_tensor* v_context = context; ggml_tensor* a_context = process_audio_context ? context : nullptr; if (is_unprocessed_dual_context) { - v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.cross_attention_dim); + v_context = ggml_ext_slice(ctx, context, 0, 0, config.cross_attention_dim); if (process_audio_context) { - a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim); + a_context = ggml_ext_slice(ctx, context, 0, config.cross_attention_dim, config.cross_attention_dim + config.audio_cross_attention_dim); } } else if (context->ne[0] == config.caption_channels * 2) { - v_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, 0, config.caption_channels); + v_context = ggml_ext_slice(ctx, context, 0, 0, config.caption_channels); if (process_audio_context) { - a_context = ggml_ext_slice(ctx->ggml_ctx, context, 0, config.caption_channels, config.caption_channels * 2); + a_context = ggml_ext_slice(ctx, context, 0, config.caption_channels, config.caption_channels * 2); } } @@ -1604,7 +1604,7 @@ namespace LTXV { auto t = ggml_repeat(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, embedded_timestep, dim, 1, embedded_timestep->ne[1]), temp); auto s = ggml_repeat(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, table, dim, 2, 1), temp); auto out = ggml_add(ctx->ggml_ctx, s, t); - return ggml_ext_chunk(ctx->ggml_ctx, out, 2, 1); + return ggml_ext_chunk(ctx, out, 2, 1); } std::pair forward(GGMLRunnerContext* ctx, @@ -1653,13 +1653,13 @@ namespace LTXV { a_context = ggml_cont(ctx->ggml_ctx, a_context); } - auto v_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, timestep, config.timestep_scale_multiplier); + auto v_timestep_scaled = ggml_ext_scale(ctx, timestep, config.timestep_scale_multiplier); auto v_pair = adaln_single->forward(ctx, v_timestep_scaled); auto v_timestep_mod = v_pair.first; auto v_embedded_time = v_pair.second; ggml_tensor* effective_audio_timestep = audio_timestep != nullptr ? audio_timestep : timestep; - auto a_timestep_scaled = ggml_ext_scale(ctx->ggml_ctx, effective_audio_timestep, config.timestep_scale_multiplier); + auto a_timestep_scaled = ggml_ext_scale(ctx, effective_audio_timestep, config.timestep_scale_multiplier); auto a_pair = audio_adaln_single->forward(ctx, a_timestep_scaled); auto a_timestep_mod = a_pair.first; auto a_embedded_time = a_pair.second; @@ -1680,13 +1680,13 @@ namespace LTXV { std::dynamic_pointer_cast(blocks["av_ca_video_scale_shift_adaln_single"])->forward(ctx, av_ca_video_timestep).first; auto av_ca_a2v_gate_noise_timestep = std::dynamic_pointer_cast(blocks["av_ca_a2v_gate_adaln_single"]) - ->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_video_timestep, av_ca_factor)) + ->forward(ctx, ggml_ext_scale(ctx, av_ca_video_timestep, av_ca_factor)) .first; auto av_ca_audio_scale_shift_timestep = std::dynamic_pointer_cast(blocks["av_ca_audio_scale_shift_adaln_single"])->forward(ctx, av_ca_audio_timestep).first; auto av_ca_v2a_gate_noise_timestep = std::dynamic_pointer_cast(blocks["av_ca_v2a_gate_adaln_single"]) - ->forward(ctx, ggml_ext_scale(ctx->ggml_ctx, av_ca_audio_timestep, av_ca_factor)) + ->forward(ctx, ggml_ext_scale(ctx, av_ca_audio_timestep, av_ca_factor)) .first; sd::ggml_graph_cut::mark_graph_cut(vx, "ltxav.prelude", "vx"); @@ -1797,9 +1797,8 @@ namespace LTXV { return {vx, ax}; } - ggml_tensor* merge_av_latents(ggml_context* ctx, - ggml_tensor* vx, - ggml_tensor* ax) const { + ggml_tensor* merge_av_latents(GGMLRunnerContext* runner_ctx, ggml_tensor* vx, ggml_tensor* ax) const { + auto ctx = runner_ctx->ggml_ctx; if (ax == nullptr || ggml_nelements(ax) == 0 || ax->ne[1] == 0) { return vx; } @@ -1815,7 +1814,7 @@ namespace LTXV { ax = ggml_cont(ctx, ax); ax = ggml_reshape_4d(ctx, ax, audio_values, 1, 1, 1); if (pad_values > 0) { - ax = ggml_ext_pad(ctx, ax, static_cast(pad_values), 0, 0, 0); + ax = ggml_ext_pad(runner_ctx, ax, static_cast(pad_values), 0, 0, 0); } int64_t extra_channels = padded_len / divisor; ax = ggml_reshape_4d(ctx, ax, width, height, frames, extra_channels); @@ -2008,7 +2007,7 @@ namespace LTXV { audio_cross_pe, video_connector_pe, audio_connector_pe); - auto out = merge_av_latents(compute_ctx, out_pair.first, out_pair.second); + auto out = merge_av_latents(&runner_ctx, out_pair.first, out_pair.second); ggml_build_forward_expand(gf, out); return gf; } diff --git a/src/model/diffusion/ming_image.hpp b/src/model/diffusion/ming_image.hpp index 1de9ac157e..8769cbe1e4 100644 --- a/src/model/diffusion/ming_image.hpp +++ b/src/model/diffusion/ming_image.hpp @@ -50,8 +50,8 @@ namespace MingImage { txt = ggml_concat(gctx, txt, direct, 1); auto t = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, timestep); const int64_t n_txt = txt->ne[1], n_img = img->ne[1]; - auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, n_txt); - auto img_pe = ggml_ext_slice(gctx, pe, 3, n_txt, n_txt + n_img); + auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, n_txt); + auto img_pe = ggml_ext_slice(ctx, pe, 3, n_txt, n_txt + n_img); for (int64_t i = 0; i < config.num_refiner_layers; ++i) { txt = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe); sd::ggml_graph_cut::mark_graph_cut(txt, "ming_image.context_refiner." + std::to_string(i), "txt"); @@ -65,9 +65,9 @@ namespace MingImage { combined = std::dynamic_pointer_cast(blocks["layers." + std::to_string(i)])->forward(ctx, combined, pe, nullptr, t); sd::ggml_graph_cut::mark_graph_cut(combined, "ming_image.layers." + std::to_string(i), "combined"); } - img = ggml_ext_slice(gctx, combined, 1, n_txt, n_txt + n_img); + img = ggml_ext_slice(ctx, combined, 1, n_txt, n_txt + n_img); img = std::dynamic_pointer_cast(blocks["final_layer"])->forward(ctx, img, t); - img = DiT::unpatchify_and_crop(gctx, img, height, width, config.patch_size, config.patch_size, false); + img = DiT::unpatchify_and_crop(ctx, img, height, width, config.patch_size, config.patch_size, false); return ggml_scale(gctx, img, -1.f); } }; diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index 9a8d91388c..ea9f9c5aea 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -263,13 +263,13 @@ namespace MiniMaxH3 { 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); + auto input = ggml_ext_slice(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)); } - return ggml_ext_vec_concat(ctx->ggml_ctx, chunks, 1); + return ggml_ext_vec_concat(ctx, chunks, 1); } }; @@ -278,29 +278,8 @@ namespace MiniMaxH3 { return ggml_reshape_3d(ctx, x, x->ne[0], x->ne[1], x->ne[2] * x->ne[3]); } - static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx, - ggml_tensor* x, - ggml_tensor* pe) { - 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, - false); - if (rot_dim == x->ne[0]) { - return rotated; - } - auto tail = attention_layout(ctx, ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0])); - return ggml_concat(ctx, rotated, tail, 0); + static ggml_tensor* apply_partial_rope(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* pe) { + return ggml_ext_rope_apply(runner_ctx, x, pe, false); } struct Attention : public GGMLBlock { @@ -381,7 +360,7 @@ namespace MiniMaxH3 { for (size_t i = 0; i < inputs.size(); ++i) { auto input = inputs[i]; for (int64_t start = 0; start < input->ne[1]; start += chunk_size) { - auto part = ggml_ext_slice(ctx->ggml_ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false); + auto part = ggml_ext_slice(ctx, input, 1, start, std::min(start + chunk_size, input->ne[1]), false); if (!ggml_is_contiguous(part)) { part = ggml_cont(ctx->ggml_ctx, part); } @@ -397,7 +376,7 @@ namespace MiniMaxH3 { } std::vector result; for (auto& parts : chunks) { - result.push_back(ggml_ext_vec_concat(ctx->ggml_ctx, parts, 2)); + result.push_back(ggml_ext_vec_concat(ctx, parts, 2)); } return result; } @@ -420,13 +399,13 @@ namespace MiniMaxH3 { for (int64_t start = 0; start < heads; start += head_chunk) { std::vector part; for (auto tensor : qkv) { - part.push_back(ggml_ext_slice(ctx->ggml_ctx, tensor, 1, start, std::min(start + head_chunk, heads), false)); + part.push_back(ggml_ext_slice(ctx, tensor, 1, start, std::min(start + head_chunk, heads), false)); } outputs.push_back(attend_heads(ctx, part, pe)); // Keep each head group's temporaries out of the concatenation tree's live set. ctx->expand_graph(outputs.back()); } - return ggml_ext_vec_concat(ctx->ggml_ctx, outputs, 0); + return ggml_ext_vec_concat(ctx, outputs, 0); } return attend_heads(ctx, qkv, pe); } @@ -562,19 +541,20 @@ namespace MiniMaxH3 { int32_t source_index = -1; }; - static std::vector modulation_row(ggml_context* ctx, + static std::vector modulation_row(GGMLRunnerContext* runner_ctx, ggml_tensor* projection, int64_t hidden_size, int expand, int modalities, int64_t row) { + auto ctx = runner_ctx->ggml_ctx; int64_t timestep_rows = projection->ne[1]; auto reshaped = ggml_reshape_2d(ctx, projection, hidden_size * expand, timestep_rows * modalities); - auto selected = ggml_ext_slice(ctx, reshaped, 1, row, row + 1); - return ggml_ext_chunk(ctx, selected, expand, 0); + auto selected = ggml_ext_slice(runner_ctx, reshaped, 1, row, row + 1); + return ggml_ext_chunk(runner_ctx, selected, expand, 0); } struct TransformerBlock : public GGMLBlock { @@ -616,9 +596,9 @@ namespace MiniMaxH3 { std::vector> rows; std::vector inputs; for (const auto& segment : segments) { - rows.push_back(modulation_row(ctx->ggml_ctx, mods, config.hidden_size, + rows.push_back(modulation_row(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)); + inputs.push_back(ggml_ext_slice(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]); @@ -629,7 +609,7 @@ namespace MiniMaxH3 { 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); + auto part = ggml_ext_slice(ctx, attended, 1, segment.start, segment.end, false); projections.push_back(attn->project_output(ctx, part)); } ggml_tensor* output = nullptr; @@ -672,8 +652,8 @@ namespace MiniMaxH3 { auto audio_out = std::dynamic_pointer_cast(blocks["audio_out"]); auto mods = adaln->forward(ctx, t_emb); auto apply = [&](const TokenModulationSpan& segment) { - 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)); + auto row = modulation_row(ctx, mods, config.hidden_size, 2, 1, segment.modulation_row); + auto value = norm->forward(ctx, ggml_ext_slice(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]); @@ -774,7 +754,7 @@ namespace MiniMaxH3 { 1); ggml_tensor* angles = nullptr; for (int axis = 0; axis < 3; ++axis) { - auto pos = ggml_ext_slice(ctx->ggml_ctx, position_ids, 0, axis, axis + 1); + auto pos = ggml_ext_slice(ctx, position_ids, 0, axis, axis + 1); auto expanded_inv = ggml_repeat_4d(ctx->ggml_ctx, inv, inv->ne[0], @@ -816,7 +796,7 @@ namespace MiniMaxH3 { ggml_tensor* video_rows = nullptr; int64_t video_offset = 0; for (auto condition : condition_videos) { - auto rows = DiT::patchify_3d(ctx->ggml_ctx, + auto rows = DiT::patchify_3d(ctx, condition, config.patch_t, config.patch_h, @@ -827,7 +807,7 @@ namespace MiniMaxH3 { video_offset += rows->ne[1]; video_rows = video_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, video_rows, rows, 1); } - auto target_video_rows = DiT::patchify_3d(ctx->ggml_ctx, + auto target_video_rows = DiT::patchify_3d(ctx, video, config.patch_t, config.patch_h, @@ -840,7 +820,7 @@ namespace MiniMaxH3 { auto video_embeds = video_proj->forward(ctx, video_rows); auto pack_audio_rows = [&](ggml_tensor* value) { - value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, value, 2, 0, 1, 3)); + value = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, value, 2, 0, 1, 3)); return ggml_reshape_3d(ctx->ggml_ctx, value, value->ne[0], @@ -873,20 +853,20 @@ namespace MiniMaxH3 { GGML_ASSERT(sequence.source_index >= 0 && sequence.source_index < static_cast(condition_video_ranges.size())); auto range = condition_video_ranges[static_cast(sequence.source_index)]; - append(ggml_ext_slice(ctx->ggml_ctx, video_embeds, 1, range.first, range.second)); + append(ggml_ext_slice(ctx, video_embeds, 1, range.first, range.second)); } else if (sequence.kind == SequenceKind::CONDITION_AUDIO) { GGML_ASSERT(sequence.source_index >= 0 && sequence.source_index < static_cast(condition_audio_ranges.size())); auto range = condition_audio_ranges[static_cast(sequence.source_index)]; - append(ggml_ext_slice(ctx->ggml_ctx, audio_embeds, 1, range.first, range.second)); + append(ggml_ext_slice(ctx, audio_embeds, 1, range.first, range.second)); } else if (sequence.kind == SequenceKind::TARGET_AUDIO) { - append(ggml_ext_slice(ctx->ggml_ctx, + append(ggml_ext_slice(ctx, audio_embeds, 1, target_audio_range.first, target_audio_range.second)); } else { - append(ggml_ext_slice(ctx->ggml_ctx, + append(ggml_ext_slice(ctx, video_embeds, 1, target_video_range.first, @@ -912,7 +892,7 @@ namespace MiniMaxH3 { auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); auto output = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment); - auto video_out = DiT::unpatchify_3d(ctx->ggml_ctx, + auto video_out = DiT::unpatchify_3d(ctx, output.first, video->ne[2] / config.patch_t, video->ne[1] / config.patch_h, @@ -927,9 +907,9 @@ namespace MiniMaxH3 { audio->ne[1] / 2, 2, audio->ne[2]); - audio_out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, audio_out, 1, 2, 0, 3)); - video_out = ggml_ext_scale(ctx->ggml_ctx, video_out, -1.f); - audio_out = ggml_ext_scale(ctx->ggml_ctx, audio_out, -1.f); + audio_out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, audio_out, 1, 2, 0, 3)); + video_out = ggml_ext_scale(ctx, video_out, -1.f); + audio_out = ggml_ext_scale(ctx, audio_out, -1.f); return {video_out, audio_out}; } }; @@ -1221,15 +1201,14 @@ namespace MiniMaxH3 { model.get_param_tensors(tensors, prefix); } - ggml_tensor* merge_av_latents(ggml_context* ctx, - ggml_tensor* video, - ggml_tensor* audio) const { + ggml_tensor* merge_av_latents(GGMLRunnerContext* runner_ctx, ggml_tensor* video, ggml_tensor* audio) const { + auto ctx = runner_ctx->ggml_ctx; int64_t divisor = video->ne[0] * video->ne[1] * video->ne[2]; int64_t values = ggml_nelements(audio); int64_t padding = (divisor - values % divisor) % divisor; audio = ggml_reshape_4d(ctx, ggml_cont(ctx, audio), values, 1, 1, 1); if (padding > 0) { - audio = ggml_ext_pad(ctx, audio, static_cast(padding), 0, 0, 0); + audio = ggml_ext_pad(runner_ctx, audio, static_cast(padding), 0, 0, 0); } audio = ggml_reshape_4d(ctx, audio, @@ -1302,21 +1281,22 @@ namespace MiniMaxH3 { } // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed // 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], - packed.shape()[2], - packed.shape()[1], - packed.shape()[0], - audio_length, - condition_videos, - condition_audios, - keyframe_indices, - reference_blocks, - text_tags, - t_v, - t_a, - r_v, - r_a); + auto audio_ctx = get_context(); + auto audio = ggml_ext_scale(&audio_ctx, audio_carrier, sigma_a / sigma_v); + auto layout = build_layout(context_tensor.shape()[1], + packed.shape()[2], + packed.shape()[1], + packed.shape()[0], + audio_length, + condition_videos, + condition_audios, + keyframe_indices, + reference_blocks, + text_tags, + t_v, + t_a, + r_v, + r_a); position_input_cache = sd::Tensor( {3, static_cast(layout.positions.size() / 3)}, @@ -1393,11 +1373,11 @@ namespace MiniMaxH3 { layout.audio_segment); // Convert the model's audio velocity to d(c_a) / d(sigma_v). output.second = ggml_add(compute_ctx, - ggml_ext_scale(compute_ctx, audio, 1.f - audio_scale), - ggml_ext_scale(compute_ctx, + ggml_ext_scale(&runner_ctx, audio, 1.f - audio_scale), + ggml_ext_scale(&runner_ctx, output.second, 1.f + (audio_scale - 1.f) * sigma_a)); - auto merged = merge_av_latents(compute_ctx, output.first, output.second); + auto merged = merge_av_latents(&runner_ctx, output.first, output.second); ggml_build_forward_expand(graph, merged); return graph; } diff --git a/src/model/diffusion/minit2i.hpp b/src/model/diffusion/minit2i.hpp index 5029759d33..63360dc837 100644 --- a/src/model/diffusion/minit2i.hpp +++ b/src/model/diffusion/minit2i.hpp @@ -257,14 +257,18 @@ namespace MiniT2I { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) { auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); - auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10000, 1.0f); + auto t_emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10000, 1.0f); t_emb = mlp_0->forward(ctx, t_emb); t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb); return mlp_2->forward(ctx, t_emb); } }; - inline std::vector split_qkv(ggml_context* ctx, ggml_tensor* qkv, int64_t num_heads, int64_t head_dim) { + inline std::vector split_qkv(GGMLRunnerContext* runner_ctx, + ggml_tensor* qkv, + int64_t num_heads, + int64_t head_dim) { + auto ctx = runner_ctx->ggml_ctx; int64_t N = qkv->ne[2]; int64_t L = qkv->ne[1]; auto q = ggml_view_4d(ctx, qkv, head_dim, num_heads, L, N, @@ -301,7 +305,7 @@ namespace MiniT2I { auto q_norm = std::dynamic_pointer_cast(blocks["q_norm"]); auto k_norm = std::dynamic_pointer_cast(blocks["k_norm"]); - auto qkv = split_qkv(ctx->ggml_ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim); + auto qkv = split_qkv(ctx, qkv_proj->forward(ctx, norm1->forward(ctx, txt)), num_heads, head_dim); auto q = q_norm->forward(ctx, qkv[0]); auto k = k_norm->forward(ctx, qkv[1]); auto v = qkv[2]; @@ -353,16 +357,16 @@ namespace MiniT2I { int64_t li = img->ne[1]; int64_t lt = txt->ne[1]; - auto img_qkv = split_qkv(ctx->ggml_ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim); - auto txt_qkv = split_qkv(ctx->ggml_ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim); + auto img_qkv = split_qkv(ctx, img_qkv_p->forward(ctx, img_norm1->forward(ctx, img)), num_heads, head_dim); + auto txt_qkv = split_qkv(ctx, txt_qkv_p->forward(ctx, txt_norm1->forward(ctx, txt)), num_heads, head_dim); auto q = ggml_concat(ctx->ggml_ctx, q_norm->forward(ctx, txt_qkv[0]), q_norm->forward(ctx, img_qkv[0]), 2); auto k = ggml_concat(ctx->ggml_ctx, k_norm->forward(ctx, txt_qkv[1]), k_norm->forward(ctx, img_qkv[1]), 2); auto v = ggml_concat(ctx->ggml_ctx, txt_qkv[2], img_qkv[2], 2); auto out = Rope::attention(ctx, q, k, v, pe, nullptr, 1.0f, false); - auto out_txt = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, lt); - auto out_img = ggml_ext_slice(ctx->ggml_ctx, out, 1, lt, lt + li); + auto out_txt = ggml_ext_slice(ctx, out, 1, 0, lt); + auto out_img = ggml_ext_slice(ctx, out, 1, lt, lt + li); img = ggml_add(ctx->ggml_ctx, img, img_proj->forward(ctx, out_img)); txt = ggml_add(ctx->ggml_ctx, txt, txt_proj->forward(ctx, out_txt)); @@ -416,7 +420,7 @@ namespace MiniT2I { mask = ggml_reshape_3d(ctx->ggml_ctx, mask, 1, mask->ne[0], mask->ne[1]); mask = ggml_repeat(ctx->ggml_ctx, mask, context); auto keep = ggml_mul(ctx->ggml_ctx, context, mask); - auto inv = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx->ggml_ctx, mask), mask); + auto inv = ggml_sub(ctx->ggml_ctx, ggml_ext_ones_like(ctx, mask), mask); auto mask_token = ggml_repeat(ctx->ggml_ctx, params["mask_token"], context); return ggml_add(ctx->ggml_ctx, keep, ggml_mul(ctx->ggml_ctx, mask_token, inv)); } @@ -468,8 +472,8 @@ namespace MiniT2I { } auto combined = ggml_concat(ctx->ggml_ctx, txt, x, 1); auto out = final_layer->forward(ctx, combined); - auto img_out = ggml_ext_slice(ctx->ggml_ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]); - return DiT::unpatchify(ctx->ggml_ctx, img_out, hp, wp, static_cast(config.patch_size), static_cast(config.patch_size), false); + auto img_out = ggml_ext_slice(ctx, out, 1, txt->ne[1], txt->ne[1] + x->ne[1]); + return DiT::unpatchify(ctx, img_out, hp, wp, static_cast(config.patch_size), static_cast(config.patch_size), false); } }; diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index b032273db5..e409a01f0a 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -272,7 +272,7 @@ struct TimestepEmbedder : public GGMLBlock { auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); - auto t_freq = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size); // [N, frequency_embedding_size] + auto t_freq = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size); // [N, frequency_embedding_size] auto t_emb = mlp_0->forward(ctx, t_freq); t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb); @@ -334,7 +334,7 @@ class SelfAttention : public GGMLBlock { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); int64_t head_dim = qkv_vec[0]->ne[0] / num_heads; auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); // [N, n_token, n_head, d_head] auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]); // [N, n_token, n_head, d_head] @@ -439,7 +439,7 @@ struct DismantledBlock : public GGMLBlock { int n_mods = 9; auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, n_mods * hidden_size] - auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0); + auto m_vec = ggml_ext_chunk(ctx, m, n_mods, 0); auto shift_msa = m_vec[0]; // [N, hidden_size] auto scale_msa = m_vec[1]; // [N, hidden_size] @@ -476,7 +476,7 @@ struct DismantledBlock : public GGMLBlock { n_mods = 2; } auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, n_mods * hidden_size] - auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, n_mods, 0); + auto m_vec = ggml_ext_chunk(ctx, m, n_mods, 0); auto shift_msa = m_vec[0]; // [N, hidden_size] auto scale_msa = m_vec[1]; // [N, hidden_size] @@ -750,7 +750,7 @@ struct FinalLayer : public GGMLBlock { auto adaLN_modulation_1 = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"]); auto m = adaLN_modulation_1->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); // [N, 2 * hidden_size] - auto m_vec = ggml_ext_chunk(ctx->ggml_ctx, m, 2, 0); + auto m_vec = ggml_ext_chunk(ctx, m, 2, 0); auto shift = m_vec[0]; // [N, hidden_size] auto scale = m_vec[1]; // [N, hidden_size] @@ -926,7 +926,7 @@ struct MMDiT : public GGMLBlock { x = forward_core_with_concat(ctx, x, c, context, skip_layers); // (N, H*W, patch_size ** 2 * out_channels) - x = DiT::unpatchify_and_crop(ctx->ggml_ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false); // [N, C, H, W] + x = DiT::unpatchify_and_crop(ctx, x, H, W, config.patch_size, config.patch_size, /*patch_last*/ false); // [N, C, H, W] return x; } diff --git a/src/model/diffusion/moss_image.hpp b/src/model/diffusion/moss_image.hpp index 2641417abf..bcca79b89f 100644 --- a/src/model/diffusion/moss_image.hpp +++ b/src/model/diffusion/moss_image.hpp @@ -164,9 +164,9 @@ namespace MOSSImage { auto c = ctx->ggml_ctx; auto linear = [&](const char* name, ggml_tensor* h) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, h); }; int64_t image_len = image->ne[0] * image->ne[1]; - auto x = linear("video_embedder", DiT::patchify(c, image, 1, 1)); + auto x = linear("video_embedder", DiT::patchify(ctx, image, 1, 1)); text = linear("text_embedder.proj", text); - auto text_pe = ggml_ext_slice(c, pe, 3, image_len, image_len + text->ne[1]); + auto text_pe = ggml_ext_slice(ctx, pe, 3, image_len, image_len + text->ne[1]); for (int i = 0; i < 2; ++i) { text = std::dynamic_pointer_cast(blocks["text_embedder.blocks." + std::to_string(i)])->forward(ctx, text, text_pe); } @@ -175,7 +175,7 @@ namespace MOSSImage { auto head = linear("video_head_modulation", time); 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); + parts[i] = ggml_ext_slice(ctx, mod, 0, i * config.dim, (i + 1) * config.dim); if (ctx->graph != nullptr) { ggml_build_forward_expand(ctx->graph, parts[i]); } @@ -187,13 +187,13 @@ namespace MOSSImage { sd::ggml_graph_cut::mark_graph_cut(x, "moss_image.layers." + std::to_string(i), "x"); } } - x = ggml_ext_slice(c, x, 1, 0, image_len); + x = ggml_ext_slice(ctx, x, 1, 0, image_len); 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); + auto shift = ggml_ext_slice(ctx, head, 0, 0, config.dim); + auto scale = ggml_ext_slice(ctx, head, 0, config.dim, 2 * config.dim); 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); + return DiT::unpatchify(ctx, x, image->ne[1], image->ne[0], 1, 1); } }; diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 0e502371a8..225a30aeb8 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -185,14 +185,15 @@ namespace Pid { return ggml_add(ctx, ggml_add(ctx, x, ggml_mul(ctx, x, scale)), shift); } - inline ggml_tensor* replicate_pad_2d(ggml_context* ctx, ggml_tensor* x) { - auto left = ggml_ext_slice(ctx, x, 0, 0, 1); - auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]); + inline ggml_tensor* replicate_pad_2d(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + auto ctx = runner_ctx->ggml_ctx; + auto left = ggml_ext_slice(runner_ctx, x, 0, 0, 1); + auto right = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]); x = ggml_concat(ctx, left, x, 0); x = ggml_concat(ctx, x, right, 0); - auto top = ggml_ext_slice(ctx, x, 1, 0, 1); - auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]); + auto top = ggml_ext_slice(runner_ctx, x, 1, 0, 1); + auto bottom = ggml_ext_slice(runner_ctx, x, 1, x->ne[1] - 1, x->ne[1]); x = ggml_concat(ctx, top, x, 1); return ggml_concat(ctx, x, bottom, 1); } @@ -235,7 +236,7 @@ namespace Pid { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* t) { auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); - auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, t, frequency_embedding_size, 10); + auto t_emb = ggml_ext_timestep_embedding(ctx, t, frequency_embedding_size, 10); t_emb = mlp_0->forward(ctx, t_emb); t_emb = ggml_silu_inplace(ctx->ggml_ctx, t_emb); return mlp_2->forward(ctx, t_emb); @@ -292,7 +293,7 @@ namespace Pid { auto proj = std::dynamic_pointer_cast(blocks["proj"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); int64_t L = x->ne[1]; int64_t N = x->ne[2]; int64_t head_dim = dim / num_heads; @@ -342,14 +343,14 @@ namespace Pid { int64_t N = x->ne[2]; int64_t head_dim = dim / num_heads; - auto qkv_x = split_qkv(ctx->ggml_ctx, qkv_x_proj->forward(ctx, x)); + auto qkv_x = ggml_ext_split_qkv(ctx, qkv_x_proj->forward(ctx, x)); auto qx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N); auto kx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N); auto vx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N); qx = q_norm_x->forward(ctx, qx); kx = k_norm_x->forward(ctx, kx); - auto qkv_y = split_qkv(ctx->ggml_ctx, qkv_y_proj->forward(ctx, y)); + auto qkv_y = ggml_ext_split_qkv(ctx, qkv_y_proj->forward(ctx, y)); auto qy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N); auto ky = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N); auto vy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N); @@ -362,8 +363,8 @@ namespace Pid { auto pos_joint = ggml_concat(ctx->ggml_ctx, pos_txt, pos_img, 3); auto out = Rope::attention(ctx, q_joint, k_joint, v_joint, pos_joint, nullptr, 1.0f, true); - auto out_y = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, Ny); - auto out_x = ggml_ext_slice(ctx->ggml_ctx, out, 1, Ny, Ny + Nx); + auto out_y = ggml_ext_slice(ctx, out, 1, 0, Ny); + auto out_x = ggml_ext_slice(ctx, out, 1, Ny, Ny + Nx); return {proj_x->forward(ctx, out_x), proj_y->forward(ctx, out_y)}; } }; @@ -400,8 +401,8 @@ namespace Pid { auto ada_img = std::dynamic_pointer_cast(blocks["adaLN_modulation_img.0"]); auto ada_txt = std::dynamic_pointer_cast(blocks["adaLN_modulation_txt.0"]); - auto mx = ggml_ext_chunk(ctx->ggml_ctx, ada_img->forward(ctx, c), 6, 0); - auto my = ggml_ext_chunk(ctx->ggml_ctx, ada_txt->forward(ctx, c), 6, 0); + auto mx = ggml_ext_chunk(ctx, ada_img->forward(ctx, c), 6, 0); + auto my = ggml_ext_chunk(ctx, ada_txt->forward(ctx, c), 6, 0); auto x_norm = apply_adaln(ctx->ggml_ctx, norm_x1->forward(ctx, x), mx[0], mx[1]); auto y_norm = apply_adaln(ctx->ggml_ctx, norm_y1->forward(ctx, y), my[0], my[1]); @@ -438,13 +439,13 @@ namespace Pid { int64_t L = (W / patch_size) * (H / patch_size); int64_t P2 = patch_size * patch_size; - auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, inputs, 2, 0, 1, 3)); + auto x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, inputs, 2, 0, 1, 3)); x = ggml_reshape_3d(ctx->ggml_ctx, x, in_channels, W * H, B); x = proj->forward(ctx, x); x = ggml_add(ctx->ggml_ctx, x, pos_full); x = ggml_reshape_4d(ctx->ggml_ctx, x, hidden_size_output, W, H, B); - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); - x = DiT::patchify(ctx->ggml_ctx, x, static_cast(patch_size), static_cast(patch_size), false); + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); + x = DiT::patchify(ctx, x, static_cast(patch_size), static_cast(patch_size), false); x = ggml_reshape_3d(ctx->ggml_ctx, x, hidden_size_output, P2, L * B); return x; } @@ -500,7 +501,7 @@ namespace Pid { auto ada_params = ada->forward(ctx, s_cond); ada_params = ggml_reshape_3d(ctx->ggml_ctx, ada_params, 6 * pixel_dim, P2, BL); - auto mod = ggml_ext_chunk(ctx->ggml_ctx, ada_params, 6, 0); + auto mod = ggml_ext_chunk(ctx, ada_params, 6, 0); auto x_norm = apply_adaln(ctx->ggml_ctx, norm1->forward(ctx, x), mod[0], mod[1]); auto x_flat = ggml_reshape_2d(ctx->ggml_ctx, x_norm, P2 * pixel_dim, BL); @@ -564,12 +565,12 @@ namespace Pid { auto conv2 = std::dynamic_pointer_cast(blocks["block.5"]); auto h = ggml_silu_inplace(ctx->ggml_ctx, norm1->forward(ctx, x)); if (replicate_padding) { - h = replicate_pad_2d(ctx->ggml_ctx, h); + h = replicate_pad_2d(ctx, h); } h = conv1->forward(ctx, h); h = ggml_silu_inplace(ctx->ggml_ctx, norm2->forward(ctx, h)); if (replicate_padding) { - h = replicate_pad_2d(ctx->ggml_ctx, h); + h = replicate_pad_2d(ctx, h); } h = conv2->forward(ctx, h); return ggml_add(ctx->ggml_ctx, x, h); @@ -628,13 +629,13 @@ namespace Pid { if (unpatchify_factor > 1) { int64_t latent_h = lq_latent->ne[1]; int64_t latent_w = lq_latent->ne[0]; - lq_latent = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, lq_latent, 2, 0, 1, 3)); + lq_latent = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, lq_latent, 2, 0, 1, 3)); lq_latent = ggml_reshape_3d(ctx->ggml_ctx, lq_latent, lq_latent->ne[0], lq_latent->ne[1] * lq_latent->ne[2], lq_latent->ne[3]); - lq_latent = DiT::unpatchify(ctx->ggml_ctx, + lq_latent = DiT::unpatchify(ctx, lq_latent, latent_h, latent_w, @@ -658,12 +659,12 @@ namespace Pid { } if (config.lq_replicate_padding) { - lq_latent = replicate_pad_2d(ctx->ggml_ctx, lq_latent); + lq_latent = replicate_pad_2d(ctx, lq_latent); } auto feat = conv0->forward(ctx, lq_latent); feat = ggml_silu_inplace(ctx->ggml_ctx, feat); if (config.lq_replicate_padding) { - feat = replicate_pad_2d(ctx->ggml_ctx, feat); + feat = replicate_pad_2d(ctx, feat); } feat = conv2->forward(ctx, feat); for (int i = 0; i < config.lq_num_res_blocks; ++i) { @@ -674,7 +675,7 @@ namespace Pid { int64_t B = feat->ne[3]; int64_t C = feat->ne[2]; int64_t L = target_pH * target_pW; - auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, feat, 2, 0, 1, 3)); + auto tokens = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, feat, 2, 0, 1, 3)); tokens = ggml_reshape_3d(ctx->ggml_ctx, tokens, C, L, B); int num_outputs = static_cast((config.patch_depth + config.lq_interval - 1) / config.lq_interval); @@ -754,15 +755,15 @@ namespace Pid { int64_t L = Hs * Ws; int64_t P2 = config.patch_size * config.patch_size; - auto x_patches = DiT::patchify(ctx->ggml_ctx, x, static_cast(config.patch_size), static_cast(config.patch_size), true); + auto x_patches = DiT::patchify(ctx, x, static_cast(config.patch_size), static_cast(config.patch_size), true); auto t_emb = t_embedder->forward(ctx, timesteps); auto condition = ggml_silu(ctx->ggml_ctx, t_emb); GGML_ASSERT(context != nullptr); int64_t Ltxt = std::min(context->ne[1], config.txt_max_length); - auto y = ggml_ext_slice(ctx->ggml_ctx, context, 1, 0, Ltxt); + auto y = ggml_ext_slice(ctx, context, 1, 0, Ltxt); auto y_emb = y_embedder->forward(ctx, y); - auto y_pos = ggml_ext_slice(ctx->ggml_ctx, params["y_pos_embedding"], 1, 0, Ltxt); + auto y_pos = ggml_ext_slice(ctx, params["y_pos_embedding"], 1, 0, Ltxt); y_emb = ggml_add(ctx->ggml_ctx, y_emb, y_pos); std::vector lq_features = lq_proj->forward(ctx, lq_latent, Hs, Ws); @@ -809,15 +810,15 @@ namespace Pid { pixels = final_layer->forward(ctx, pixels); pixels = ggml_reshape_3d(ctx->ggml_ctx, pixels, config.in_channels * P2, L, B); - auto out = DiT::unpatchify(ctx->ggml_ctx, + auto out = DiT::unpatchify(ctx, pixels, Hs, Ws, static_cast(config.patch_size), static_cast(config.patch_size), false); - out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H_orig); - out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W_orig); + out = ggml_ext_slice(ctx, out, 1, 0, H_orig); + out = ggml_ext_slice(ctx, out, 0, 0, W_orig); return out; } }; diff --git a/src/model/diffusion/pixart.hpp b/src/model/diffusion/pixart.hpp index 74bde6d453..ca8aba2d7f 100644 --- a/src/model/diffusion/pixart.hpp +++ b/src/model/diffusion/pixart.hpp @@ -140,11 +140,11 @@ namespace PixArt { std::vector qkv; if (self_attention) { auto projected = std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x); - qkv = ggml_ext_chunk(ctx->ggml_ctx, projected, 3, 0); + qkv = ggml_ext_chunk(ctx, projected, 3, 0); } else { auto q = std::dynamic_pointer_cast(blocks["q_linear"])->forward(ctx, x); auto kv = std::dynamic_pointer_cast(blocks["kv_linear"])->forward(ctx, context); - auto parts = ggml_ext_chunk(ctx->ggml_ctx, kv, 2, 0); + auto parts = ggml_ext_chunk(ctx, kv, 2, 0); qkv = {q, parts[0], parts[1]}; } auto out = ggml_ext_attention_ext(ctx, qkv[0], qkv[1], qkv[2], num_heads, mask, false, ctx->flash_attn_enabled); @@ -186,7 +186,7 @@ namespace PixArt { } table = ggml_reshape_3d(ctx->ggml_ctx, table, dim, 6, 1); auto m = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, mod, dim, 6, N), table); - auto mv = ggml_ext_chunk(ctx->ggml_ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, m), dim * 6, N), 6, 0); + auto mv = ggml_ext_chunk(ctx, ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, m), dim * 6, N), 6, 0); auto attn1 = std::dynamic_pointer_cast(blocks["attn"]); auto attn2 = std::dynamic_pointer_cast(blocks["cross_attn"]); @@ -204,7 +204,7 @@ namespace PixArt { x = ggml_add(ctx->ggml_ctx, x, attn2->forward(ctx, x, context, context_mask)); 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 = ggml_ext_gelu(ctx, h, true); h = fc2->forward(ctx, h); return gate(h, mv[5]); } @@ -256,18 +256,18 @@ namespace PixArt { int64_t hp = H / p; auto h = std::dynamic_pointer_cast(blocks["x_embedder.proj"])->forward(ctx, x); // [N, hidden, hp, wp] - h = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3)); // [N, hp, wp, hidden] -> [N, hp*wp, hidden] + h = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, h, 1, 2, 0, 3)); // [N, hp, wp, hidden] -> [N, hp*wp, hidden] h = ggml_reshape_3d(ctx->ggml_ctx, h, config.hidden_size, wp * hp, N); // [N, hp*wp, hidden] h = ggml_add(ctx->ggml_ctx, h, pos_embed); - auto t = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, ADALN_EMBED_DIM, 10000); + auto t = ggml_ext_timestep_embedding(ctx, timesteps, ADALN_EMBED_DIM, 10000); auto emb = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); auto mod = std::dynamic_pointer_cast(blocks["t_block.1"]) ->forward(ctx, ggml_silu(ctx->ggml_ctx, emb)); // [N, 6 * hidden] auto ctx_emb = std::dynamic_pointer_cast(blocks["y_embedder.y_proj.fc1"])->forward(ctx, context); - ctx_emb = ggml_ext_gelu(ctx->ggml_ctx, ctx_emb, true); + ctx_emb = ggml_ext_gelu(ctx, ctx_emb, true); ctx_emb = std::dynamic_pointer_cast(blocks["y_embedder.y_proj.fc2"])->forward(ctx, ctx_emb); for (int i = 0; i < config.num_layers; ++i) { @@ -284,13 +284,13 @@ namespace PixArt { auto ss = ggml_add(ctx->ggml_ctx, ggml_reshape_3d(ctx->ggml_ctx, tail_table, config.hidden_size, 2, 1), ggml_reshape_3d(ctx->ggml_ctx, emb, config.hidden_size, 1, N)); // [2, hidden, N] - auto parts = ggml_ext_chunk(ctx->ggml_ctx, - ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, ss), config.hidden_size * 2, N), + auto parts = ggml_ext_chunk(ctx, + ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(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, 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); + h = DiT::unpatchify(ctx, h, hp, wp, static_cast(p), static_cast(p), false); return h; // [N, out_channels, H, W] } }; @@ -369,7 +369,7 @@ namespace PixArt { auto runner_ctx = get_context(); ggml_tensor* out = model.forward(&runner_ctx, x, timesteps, context, pos, context_mask); // learn_sigma: keep the noise prediction half of the output channels - out = ggml_ext_slice(compute_ctx, out, 2, 0, config.in_channels); + out = ggml_ext_slice(&runner_ctx, out, 2, 0, config.in_channels); ggml_build_forward_expand(gf, out); return gf; } diff --git a/src/model/diffusion/qwen_image.hpp b/src/model/diffusion/qwen_image.hpp index 6435222f59..a9aa41bb8c 100644 --- a/src/model/diffusion/qwen_image.hpp +++ b/src/model/diffusion/qwen_image.hpp @@ -113,7 +113,7 @@ namespace Qwen { // return: [N, embedding_dim] auto timestep_embedder = std::dynamic_pointer_cast(blocks["timestep_embedder"]); - auto timesteps_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, 256, 10000, 1.f); + auto timesteps_proj = ggml_ext_timestep_embedding(ctx, timesteps, 256, 10000, 1.f); auto timesteps_emb = timestep_embedder->forward(ctx, timesteps_proj); if (use_additional_t_cond) { GGML_ASSERT(addition_t_cond != nullptr); @@ -186,9 +186,7 @@ namespace Qwen { auto to_v = std::dynamic_pointer_cast(blocks["to_v"]); auto to_out_0 = std::dynamic_pointer_cast(blocks["to_out.0"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - to_out_0->set_force_prec_f32(true); - } + to_out_0->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); auto norm_added_q = std::dynamic_pointer_cast(blocks["norm_added_q"]); auto norm_added_k = std::dynamic_pointer_cast(blocks["norm_added_k"]); @@ -289,14 +287,17 @@ namespace Qwen { eps)); } - std::vector get_mod_params_vec(ggml_context* ctx, ggml_tensor* mod_params, ggml_tensor* index = nullptr) { + std::vector get_mod_params_vec(GGMLRunnerContext* runner_ctx, + ggml_tensor* mod_params, + ggml_tensor* index = nullptr) { + auto ctx = runner_ctx->ggml_ctx; // index: [N, n_img_token] // mod_params: [N, hidden_size * 12] if (index == nullptr) { - return ggml_ext_chunk(ctx, mod_params, 6, 0); + return ggml_ext_chunk(runner_ctx, mod_params, 6, 0); } mod_params = ggml_reshape_1d(ctx, mod_params, ggml_nelements(mod_params)); - auto mod_params_vec = ggml_ext_chunk(ctx, mod_params, 12, 0); + auto mod_params_vec = ggml_ext_chunk(runner_ctx, mod_params, 12, 0); index = ggml_reshape_3d(ctx, index, 1, index->ne[0], index->ne[1]); // [N, n_img_token, 1] index = ggml_repeat_4d(ctx, index, mod_params_vec[0]->ne[0], index->ne[1], index->ne[2], index->ne[3]); // [N, n_img_token, hidden_size] std::vector mod_results; @@ -339,15 +340,15 @@ namespace Qwen { auto img_mod_params = ggml_silu(ctx->ggml_ctx, t_emb); img_mod_params = img_mod_1->forward(ctx, img_mod_params); - auto img_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, img_mod_params, modulate_index); + auto img_mod_param_vec = get_mod_params_vec(ctx, img_mod_params, modulate_index); if (zero_cond_t) { - t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0]; + t_emb = ggml_ext_chunk(ctx, t_emb, 2, 1)[0]; } auto txt_mod_params = ggml_silu(ctx->ggml_ctx, t_emb); txt_mod_params = txt_mod_1->forward(ctx, txt_mod_params); - auto txt_mod_param_vec = get_mod_params_vec(ctx->ggml_ctx, txt_mod_params); + auto txt_mod_param_vec = get_mod_params_vec(ctx, txt_mod_params); auto img_normed = img_norm1->forward(ctx, img); auto img_modulated = Flux::modulate(ctx, img_normed, img_mod_param_vec[0], img_mod_param_vec[1], modulate_index != nullptr); @@ -402,7 +403,7 @@ namespace Qwen { auto linear = std::dynamic_pointer_cast(blocks["linear"]); auto emb = linear->forward(ctx, ggml_silu(ctx->ggml_ctx, c)); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, emb, 2, 0); + auto mods = ggml_ext_chunk(ctx, emb, 2, 0); auto scale = mods[0]; auto shift = mods[1]; @@ -461,7 +462,7 @@ namespace Qwen { auto t_emb = time_text_embed->forward(ctx, timestep, addition_t_cond); if (config.zero_cond_t) { - auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), addition_t_cond); + auto t_emb_0 = time_text_embed->forward(ctx, ggml_ext_zeros_like(ctx, timestep), addition_t_cond); t_emb = ggml_concat(ctx->ggml_ctx, t_emb, t_emb_0, 1); } auto img = img_in->forward(ctx, x); @@ -482,7 +483,7 @@ namespace Qwen { } if (config.zero_cond_t) { - t_emb = ggml_ext_chunk(ctx->ggml_ctx, t_emb, 2, 1)[0]; + t_emb = ggml_ext_chunk(ctx, t_emb, 2, 1)[0]; } img = norm_out->forward(ctx, img, t_emb); @@ -519,12 +520,12 @@ namespace Qwen { auto patchify_input = [&](ggml_tensor* input) -> ggml_tensor* { input = DiT::pad_to_patch_size(ctx, input, config.patch_size, config.patch_size); if (!has_time_axis) { - return DiT::patchify(ctx->ggml_ctx, input, config.patch_size, config.patch_size); + return DiT::patchify(ctx, input, config.patch_size, config.patch_size); } if (input->ne[3] == 1) { input = ggml_reshape_4d(ctx->ggml_ctx, input, input->ne[0], input->ne[1], 1, input->ne[2]); } - return DiT::patchify_3d(ctx->ggml_ctx, input, 1, config.patch_size, config.patch_size, N); + return DiT::patchify_3d(ctx, input, 1, config.patch_size, config.patch_size, N); }; auto img = patchify_input(x); @@ -550,11 +551,11 @@ namespace Qwen { int pad_w = (config.patch_size - W % config.patch_size) % config.patch_size; int h_len = static_cast((H + pad_h) / config.patch_size); int w_len = static_cast((W + pad_w) / config.patch_size); - out = DiT::unpatchify_3d(ctx->ggml_ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size); - out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] + out = DiT::unpatchify_3d(ctx, out, T, h_len, w_len, 1, config.patch_size, config.patch_size); + out = ggml_ext_slice(ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] + out = ggml_ext_slice(ctx, out, 0, 0, W); // [N*C, T, H, W] } else { - out = DiT::unpatchify_and_crop(ctx->ggml_ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W] + out = DiT::unpatchify_and_crop(ctx, out, H, W, config.patch_size, config.patch_size); // [N, C, H, W] } return out; diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 9452ec4678..7cd86fd3e7 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -146,7 +146,7 @@ namespace Qwen { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { auto weight = params["weight"]; if (ctx->weight_adapter) { - weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight"); + weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight"); } weight = ggml_scale_bias(ctx->ggml_ctx, weight, 1.f, 1.f); return ggml_mul(ctx->ggml_ctx, ggml_rms_norm(ctx->ggml_ctx, x, eps), weight); @@ -164,7 +164,7 @@ namespace Qwen { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { x = std::dynamic_pointer_cast(blocks["text_norm"])->forward(ctx, x); x = std::dynamic_pointer_cast(blocks["in_layer"])->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = ggml_ext_gelu(ctx, x); return std::dynamic_pointer_cast(blocks["out_layer"])->forward(ctx, x); } }; @@ -196,7 +196,7 @@ namespace Qwen { // prefix before its full-sequence K/V can accumulate across layers. ctx->expand_graph(q); auto persist = [&](ggml_tensor* tensor, int axis, const char* name) { - auto part = ggml_ext_slice(ctx->ggml_ctx, tensor, axis, 0, cache.prefix_length); + auto part = ggml_ext_slice(ctx, tensor, axis, 0, cache.prefix_length); // Pack the contiguous data into wider rows so quantization blocks // can exceed head_dim without padding or changing element order. part = ggml_reshape_2d(ctx->ggml_ctx, part, x->ne[0], cache.prefix_length); @@ -236,17 +236,15 @@ namespace Qwen { } else { for (size_t i = 0; i < segments.size(); ++i) { const auto& segment = segments[i]; - auto sq = ggml_ext_slice(ctx->ggml_ctx, q, 1, segment.start, segment.end); - auto sk = ggml_ext_slice(ctx->ggml_ctx, k, 1, 0, segment.end); - auto sv = ggml_ext_slice(ctx->ggml_ctx, v, 2, 0, segment.end); + auto sq = ggml_ext_slice(ctx, q, 1, segment.start, segment.end); + auto sk = ggml_ext_slice(ctx, k, 1, 0, segment.end); + auto sv = ggml_ext_slice(ctx, v, 2, 0, segment.end); auto out = attend(sq, sk, sv, masks[i]); result = result == nullptr ? out : ggml_concat(ctx->ggml_ctx, result, out, 1); } } auto to_out = std::dynamic_pointer_cast(blocks["to_out.0"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - to_out->set_force_prec_f32(true); - } + to_out->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); return to_out->forward(ctx, result); } }; @@ -268,7 +266,7 @@ namespace Qwen { 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 rows = ggml_ext_chunk(runner_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); @@ -277,8 +275,8 @@ namespace Qwen { if (prefix_length == 0) { 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]); + auto target = apply(ggml_ext_slice(runner_ctx, x, 1, prefix_length, x->ne[1]), rows[0]); + auto prefix = apply(ggml_ext_slice(runner_ctx, x, 1, 0, prefix_length), rows[1]); return ggml_concat(ctx, prefix, target, 1); } @@ -323,27 +321,27 @@ namespace Qwen { } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* timestep, ggml_tensor* context, const std::vector& refs, ggml_tensor* pe, const QwenImage21Layout& layout, const std::vector& masks, const QwenImage21PrefixCache& cache) { - auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx->ggml_ctx, timestep), 0); + auto time = ggml_concat(ctx->ggml_ctx, timestep, ggml_ext_zeros_like(ctx, timestep), 0); // Runtime flow timesteps already use the [0, 1000] scale. - time = ggml_ext_timestep_embedding(ctx->ggml_ctx, time, 256, 10000, 1.f); + time = ggml_ext_timestep_embedding(ctx, time, 256, 10000, 1.f); time = std::dynamic_pointer_cast(blocks["time_text_embed.timestep_embedder"])->forward(ctx, time); time = ggml_silu(ctx->ggml_ctx, time); auto modulation = std::dynamic_pointer_cast(blocks["modulation.1"])->forward(ctx, time); - auto mod = ggml_ext_chunk(ctx->ggml_ctx, modulation, 4, 0); + auto mod = ggml_ext_chunk(ctx, modulation, 4, 0); auto img_in = std::dynamic_pointer_cast(blocks["img_in"]); ggml_tensor* joint = nullptr; if (cache.mode == QwenImage21PrefixCache::Mode::REUSE) { - joint = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, x, 1, 1)); + joint = img_in->forward(ctx, DiT::patchify(ctx, x, 1, 1)); } else { auto text = std::dynamic_pointer_cast(blocks["txt_in"])->forward(ctx, context); for (const auto& segment : layout.segments) { ggml_tensor* h; if (segment.image_index < 0) { - h = ggml_ext_slice(ctx->ggml_ctx, text, 1, segment.context_start, + h = ggml_ext_slice(ctx, text, 1, segment.context_start, segment.context_start + segment.end - segment.start); } else { auto image = segment.image_index == static_cast(refs.size()) ? x : refs[segment.image_index]; - h = img_in->forward(ctx, DiT::patchify(ctx->ggml_ctx, image, 1, 1)); + h = img_in->forward(ctx, DiT::patchify(ctx, image, 1, 1)); } joint = joint == nullptr ? h : ggml_concat(ctx->ggml_ctx, joint, h, 1); } @@ -359,13 +357,13 @@ namespace Qwen { sd::ggml_graph_cut::mark_graph_cut(joint, layer_cache.cut_group, "joint"); } if (cache.mode != QwenImage21PrefixCache::Mode::REUSE) { - joint = ggml_ext_slice(ctx->ggml_ctx, joint, 1, layout.prefix_length, joint->ne[1]); + joint = ggml_ext_slice(ctx, joint, 1, layout.prefix_length, joint->ne[1]); } - auto scale = std::dynamic_pointer_cast(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx->ggml_ctx, time, 2, 1)[0]); + auto scale = std::dynamic_pointer_cast(blocks["norm_out.linear"])->forward(ctx, ggml_ext_chunk(ctx, time, 2, 1)[0]); joint = std::dynamic_pointer_cast(blocks["norm_out.norm"])->forward(ctx, joint); joint = ggml_mul(ctx->ggml_ctx, joint, ggml_scale_bias(ctx->ggml_ctx, scale, 1.f, 1.f)); joint = std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, joint); - return DiT::unpatchify_and_crop(ctx->ggml_ctx, joint, x->ne[1], x->ne[0], 1, 1); + return DiT::unpatchify_and_crop(ctx, joint, x->ne[1], x->ne[0], 1, 1); } }; diff --git a/src/model/diffusion/relay_vsr.hpp b/src/model/diffusion/relay_vsr.hpp index a390aa34f1..12c7a1ad35 100644 --- a/src/model/diffusion/relay_vsr.hpp +++ b/src/model/diffusion/relay_vsr.hpp @@ -124,7 +124,7 @@ namespace RelayVSR { k = VSR::apply_rope(ctx, k, key_rope, heads_); out = local->forward(ctx, q, k, v, heads_); } else { - out = ggml_ext_attention_ext(c, ctx->backend, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled); + out = ggml_ext_attention_ext(ctx, q, k, v, heads_, nullptr, false, ctx->flash_attn_enabled); } return to_out_0->forward(ctx, out); } @@ -157,7 +157,7 @@ namespace RelayVSR { auto ffn_net_2 = std::dynamic_pointer_cast(blocks["ffn.net.2"]); auto c = ctx->ggml_ctx; auto modulation = ggml_add(c, timestep_proj, params["scale_shift_table"]); - auto pieces = ggml_ext_chunk(c, modulation, 6, 1); + auto pieces = ggml_ext_chunk(ctx, modulation, 6, 1); auto attn1 = std::dynamic_pointer_cast(blocks["attn1"]); auto attn2 = std::dynamic_pointer_cast(blocks["attn2"]); auto y = attn1->forward(ctx, ggml_ext_modulate(ctx, norm1->forward(ctx, x), pieces[0], pieces[1]), nullptr, qpe, kpe, &local, name + ".attn1"); @@ -165,7 +165,7 @@ namespace RelayVSR { x = ggml_add(c, x, attn2->forward(ctx, norm2->forward(ctx, x), text, nullptr, nullptr, nullptr, name + ".attn2")); y = ggml_ext_modulate(ctx, norm3->forward(ctx, x), pieces[3], pieces[4]); y = ffn_net_0_proj->forward(ctx, y); - y = ffn_net_2->forward(ctx, ggml_ext_gelu(c, y)); + y = ffn_net_2->forward(ctx, ggml_ext_gelu(ctx, y)); return ggml_add(c, x, ggml_mul(c, y, pieces[5])); } }; @@ -213,7 +213,7 @@ namespace RelayVSR { auto e0 = condition_embedder_time_proj->forward(ctx, ggml_silu(c, e)); e0 = ggml_reshape_3d(c, e0, config_.dim, 6, 1); auto text = condition_embedder_text_embedder_linear_1->forward(ctx, params["prompt_embeds"]); - text = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(c, text)); + text = condition_embedder_text_embedder_linear_2->forward(ctx, ggml_ext_gelu(ctx, text)); for (int i = 0; i < config_.num_layers; ++i) { auto name = "blocks." + std::to_string(i); x = std::dynamic_pointer_cast(blocks[name])->forward(ctx, x, text, e0, qpe, kpe, local, name); @@ -221,7 +221,7 @@ namespace RelayVSR { } e = ggml_reshape_3d(c, e, config_.dim, 1, 1); auto modulation = ggml_add(c, params["scale_shift_table"], e); - auto pieces = ggml_ext_chunk(c, modulation, 2, 1); + auto pieces = ggml_ext_chunk(ctx, modulation, 2, 1); x = norm_out->forward(ctx, x); x = ggml_ext_modulate(ctx, x, pieces[0], pieces[1]); x = proj_out->forward(ctx, x); diff --git a/src/model/diffusion/seedvr2.hpp b/src/model/diffusion/seedvr2.hpp index 87e542e910..0f0a62e9d9 100644 --- a/src/model/diffusion/seedvr2.hpp +++ b/src/model/diffusion/seedvr2.hpp @@ -139,7 +139,7 @@ namespace SeedVR2 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { auto up = std::dynamic_pointer_cast(blocks["proj_in"])->forward(ctx, x); auto hidden = swiglu_ ? ggml_swiglu_split(ctx->ggml_ctx, std::dynamic_pointer_cast(blocks["proj_in_gate"])->forward(ctx, x), up) - : ggml_ext_gelu(ctx->ggml_ctx, up); + : ggml_ext_gelu(ctx, up); return std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, hidden); } }; @@ -200,7 +200,7 @@ namespace SeedVR2 { auto vpe = rotation(ctx, inputs.input(ctx, plan.positions)); auto tpe = rotation(ctx, inputs.input(ctx, plan.text_positions)); auto qkv = [&](ggml_tensor* x, bool text, ggml_tensor* pe) { - auto parts = ggml_ext_chunk(c, x, 3, 0); + auto parts = ggml_ext_chunk(ctx, x, 3, 0); for (int i = 0; i < 2; ++i) { auto h = ggml_reshape_3d(c, parts[i], 128, heads, x->ne[1]); h = std::dynamic_pointer_cast(blocks[std::string(i == 0 ? "attn.norm_q." : "attn.norm_k.") + branch(text)])->forward(ctx, h); @@ -216,15 +216,15 @@ namespace SeedVR2 { for (int length : plan.lengths) { std::vector joined; for (int i = 0; i < 3; ++i) - joined.push_back(ggml_concat(c, ggml_ext_slice(c, vp[i], 1, start, start + length), tp[i], 1)); - auto out = ggml_ext_attention_ext(c, ctx->backend, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled); - auto vo = ggml_ext_slice(c, out, 1, 0, length); - auto to = ggml_ext_slice(c, out, 1, length, length + txt->ne[1]); + joined.push_back(ggml_concat(c, ggml_ext_slice(ctx, vp[i], 1, start, start + length), tp[i], 1)); + auto out = ggml_ext_attention_ext(ctx, joined[0], joined[1], joined[2], heads, nullptr, false, ctx->flash_attn_enabled); + auto vo = ggml_ext_slice(ctx, out, 1, 0, length); + auto to = ggml_ext_slice(ctx, out, 1, length, length + txt->ne[1]); video_windows.push_back(vo); tout = tout ? ggml_add(c, tout, to) : to; start += length; } - auto vout = ggml_ext_vec_concat(c, video_windows, 1); + auto vout = ggml_ext_vec_concat(ctx, video_windows, 1); vout = ggml_get_rows(c, vout, inputs.input(ctx, plan.reverse)); tout = ggml_scale(c, tout, 1.f / plan.lengths.size()); vid = ggml_add(c, vid, ada(linear("attn.proj_out", false, vout), false, "attn", true)); diff --git a/src/model/diffusion/sefi_image.hpp b/src/model/diffusion/sefi_image.hpp index 7781532260..150c8bf1f3 100644 --- a/src/model/diffusion/sefi_image.hpp +++ b/src/model/diffusion/sefi_image.hpp @@ -77,8 +77,8 @@ namespace SefiImage { auto semantic_embedder = std::dynamic_pointer_cast(blocks["semantic_embedder"]); auto texture_embedder = std::dynamic_pointer_cast(blocks["texture_embedder"]); - auto sem_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f); - auto tex_proj = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f); + auto sem_proj = ggml_ext_timestep_embedding(ctx, timestep_sem, (int)timestep_guidance_in_dim, 10000, 1.f); + auto tex_proj = ggml_ext_timestep_embedding(ctx, timestep_tex, (int)timestep_guidance_in_dim, 10000, 1.f); auto sem_emb = semantic_embedder->forward(ctx, sem_proj); auto tex_emb = texture_embedder->forward(ctx, tex_proj); return ggml_concat(ctx->ggml_ctx, sem_emb, tex_emb, 0); diff --git a/src/model/diffusion/sensenova_u1.h b/src/model/diffusion/sensenova_u1.h index 302df6b26c..45dfbadb81 100644 --- a/src/model/diffusion/sensenova_u1.h +++ b/src/model/diffusion/sensenova_u1.h @@ -143,7 +143,7 @@ namespace SenseNovaU1 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timesteps) { auto mlp_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto mlp_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); - auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx, + auto x = ggml_ext_timestep_embedding(ctx, timesteps, static_cast(frequency_embedding_size), 10000, @@ -166,8 +166,8 @@ namespace SenseNovaU1 { // head axis used by the RoPE kernel: [hidden, 1, tokens, batch]. x = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], 1, x->ne[1], x->ne[2]); const int64_t half = x->ne[0] / 2; - auto x_part = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, half); - auto y_part = ggml_ext_slice(ctx->ggml_ctx, x, 0, half, x->ne[0]); + auto x_part = ggml_ext_slice(ctx, x, 0, 0, half); + auto y_part = ggml_ext_slice(ctx, x, 0, half, x->ne[0]); x_part = ggml_rope_ext(ctx->ggml_ctx, x_part, position_x, @@ -255,10 +255,10 @@ namespace SenseNovaU1 { const int64_t h = x->ne[1]; const int64_t w = x->ne[0]; GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0); - x = ggml_ext_cont(ctx->ggml_ctx, - ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = ggml_ext_cont(ctx, + ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); - return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true); + return DiT::unpatchify(ctx, x, h, w, upscale_factor, upscale_factor, true); } struct PixelDecoder : public GGMLBlock { @@ -352,13 +352,13 @@ namespace SenseNovaU1 { const int64_t spatial_dim = config.head_dim - temporal_dim; const int64_t axis_dim = spatial_dim / 2; - auto temporal = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, temporal_dim); - auto spatial = ggml_ext_slice(ctx->ggml_ctx, x, 0, temporal_dim, config.head_dim); + auto temporal = ggml_ext_slice(ctx, x, 0, 0, temporal_dim); + auto spatial = ggml_ext_slice(ctx, x, 0, temporal_dim, config.head_dim); temporal = std::dynamic_pointer_cast(blocks[norm_name])->forward(ctx, temporal); spatial = std::dynamic_pointer_cast(blocks[norm_hw_name])->forward(ctx, spatial); - auto height = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, 0, axis_dim); - auto width = ggml_ext_slice(ctx->ggml_ctx, spatial, 0, axis_dim, spatial_dim); + auto height = ggml_ext_slice(ctx, spatial, 0, 0, axis_dim); + auto width = ggml_ext_slice(ctx, spatial, 0, axis_dim, spatial_dim); temporal = apply_axis_rope(ctx, temporal, position_t, @@ -448,14 +448,13 @@ namespace SenseNovaU1 { } q = ggml_cont(ctx->ggml_ctx, - ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)); + ggml_ext_torch_permute(ctx, q, 0, 2, 1, 3)); q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]); k = ggml_cont(ctx->ggml_ctx, - ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); + ggml_ext_torch_permute(ctx, k, 0, 2, 1, 3)); k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]); - auto out = ggml_ext_attention_ext(ctx->ggml_ctx, - ctx->backend, + auto out = ggml_ext_attention_ext(ctx, q, k, v, diff --git a/src/model/diffusion/unet.hpp b/src/model/diffusion/unet.hpp index 837163fd90..767dfc1efa 100644 --- a/src/model/diffusion/unet.hpp +++ b/src/model/diffusion/unet.hpp @@ -247,7 +247,7 @@ class SpatialVideoTransformer : public SpatialTransformer { auto num_frames = ggml_arange(ctx->ggml_ctx, 0.f, static_cast(timesteps), 1.f); // since b is 1, no need to do repeat - auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, num_frames, static_cast(in_channels), max_time_embed_period); // [N, in_channels] + auto t_emb = ggml_ext_timestep_embedding(ctx, num_frames, static_cast(in_channels), max_time_embed_period); // [N, in_channels] auto emb = time_pos_embed_0->forward(ctx, t_emb); emb = ggml_silu_inplace(ctx->ggml_ctx, emb); @@ -572,7 +572,7 @@ class UnetModelBlock : public GGMLBlock { auto out_0 = std::dynamic_pointer_cast(blocks["out.0"]); auto out_2 = std::dynamic_pointer_cast(blocks["out.2"]); - auto t_emb = ggml_ext_timestep_embedding(ctx->ggml_ctx, timesteps, model_channels); // [N, model_channels] + auto t_emb = ggml_ext_timestep_embedding(ctx, timesteps, model_channels); // [N, model_channels] auto emb = time_embed_0->forward(ctx, t_emb); emb = ggml_silu_inplace(ctx->ggml_ctx, emb); @@ -685,7 +685,7 @@ class UnetModelBlock : public GGMLBlock { sd::ggml_graph_cut::mark_graph_cut(h, "unet.middle_block", "h"); if (controls.size() > 0) { // not in place, it segfaults the Vulkan backend with SDXL ControlNets (seen on RADV) - auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[controls.size() - 1], control_strength, false); + auto cs = ggml_ext_scale(ctx, controls[controls.size() - 1], control_strength, false); h = ggml_add(ctx->ggml_ctx, h, cs); // middle control } int control_offset = static_cast(controls.size() - 2); @@ -699,7 +699,7 @@ class UnetModelBlock : public GGMLBlock { if (controls.size() > 0) { // not in place, see the middle control - auto cs = ggml_ext_scale(ctx->ggml_ctx, controls[control_offset], control_strength, false); + auto cs = ggml_ext_scale(ctx, controls[control_offset], control_strength, false); h_skip = ggml_add(ctx->ggml_ctx, h_skip, cs); // control net condition control_offset--; } diff --git a/src/model/diffusion/wan.hpp b/src/model/diffusion/wan.hpp index 3c76c4ea19..124020e05f 100644 --- a/src/model/diffusion/wan.hpp +++ b/src/model/diffusion/wan.hpp @@ -362,7 +362,7 @@ namespace WAN { auto modulation = params["modulation"]; e = ggml_add(ctx->ggml_ctx, e, modulation); // [N, 6, dim] or [N, T, 6, dim] - auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 6, 1); // ([N, 1, dim], ...) or [N, T, 1, dim] + auto es = ggml_ext_chunk(ctx, e, 6, 1); // ([N, 1, dim], ...) or [N, T, 1, dim] auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); auto self_attn = std::dynamic_pointer_cast(blocks["self_attn"]); @@ -391,7 +391,7 @@ namespace WAN { y = modulate_add(ctx->ggml_ctx, y, es[3]); y = ffn_0->forward(ctx, y); - y = ggml_ext_gelu(ctx->ggml_ctx, y, true); + y = ggml_ext_gelu(ctx, y, true); y = ffn_2->forward(ctx, y); x = ggml_add(ctx->ggml_ctx, x, modulate_mul(ctx->ggml_ctx, y, es[5])); @@ -484,7 +484,7 @@ namespace WAN { e = ggml_repeat_4d(ctx->ggml_ctx, e, e->ne[0], 2, e->ne[2], e->ne[3]); // [N, 2, dim] or [N, T, 2, dim] e = ggml_add(ctx->ggml_ctx, e, modulation); // [N, 2, dim] or [N, T, 2, dim] - auto es = ggml_ext_chunk(ctx->ggml_ctx, e, 2, 1); // ([N, 1, dim], ...) or ([N, T, 1, dim], ...) + auto es = ggml_ext_chunk(ctx, e, 2, 1); // ([N, 1, dim], ...) or ([N, T, 1, dim], ...) auto norm = std::dynamic_pointer_cast(blocks["norm"]); auto head = std::dynamic_pointer_cast(blocks["head"]); @@ -525,8 +525,8 @@ namespace WAN { if (flf_pos_embed_token_number > 0) { auto emb_pos = params["emb_pos"]; - auto a = ggml_ext_slice(ctx->ggml_ctx, image_embeds, 1, 0, emb_pos->ne[1]); - auto b = ggml_ext_slice(ctx->ggml_ctx, emb_pos, 1, 0, image_embeds->ne[1]); + auto a = ggml_ext_slice(ctx, image_embeds, 1, 0, emb_pos->ne[1]); + auto b = ggml_ext_slice(ctx, emb_pos, 1, 0, image_embeds->ne[1]); image_embeds = ggml_add(ctx->ggml_ctx, a, b); } @@ -538,7 +538,7 @@ namespace WAN { auto x = proj_0->forward(ctx, image_embeds); x = proj_1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); x = proj_3->forward(ctx, x); x = proj_4->forward(ctx, x); @@ -639,15 +639,16 @@ namespace WAN { int pad_t = (std::get<0>(config.patch_size) - T % std::get<0>(config.patch_size)) % std::get<0>(config.patch_size); int pad_h = (std::get<1>(config.patch_size) - H % std::get<1>(config.patch_size)) % std::get<1>(config.patch_size); int pad_w = (std::get<2>(config.patch_size) - W % std::get<2>(config.patch_size)) % std::get<2>(config.patch_size); - ggml_ext_pad(ctx->ggml_ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + ggml_ext_pad(ctx, x, pad_w, pad_h, pad_t, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); return x; } - ggml_tensor* unpatchify(ggml_context* ctx, + ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t t_len, int64_t h_len, int64_t w_len) { + auto ctx = runner_ctx->ggml_ctx; // x: [N, t_len*h_len*w_len, pt*ph*pw*C] // return: [N*C, t_len*pt, h_len*ph, w_len*pw] int64_t N = x->ne[3]; @@ -658,15 +659,15 @@ namespace WAN { GGML_ASSERT(C * pt * ph * pw == x->ne[0]); - x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] - x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] - x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, C, pw * ph * pt, w_len * h_len * t_len, N); // [N, t_len*h_len*w_len, pt*ph*pw, C] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3)); // [N, C, t_len*h_len*w_len, pt*ph*pw] + x = ggml_reshape_4d(ctx, x, pw, ph * pt, w_len, h_len * t_len * C * N); // [N*C*t_len*h_len, w_len, pt*ph, pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, pt*ph, w_len, pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph, pt, h_len * t_len * C * N); // [N*C*t_len*h_len, pt, ph, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len*h_len, ph, pt, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, pt, ph * h_len, t_len * C * N); // [N*C*t_len, h_len*ph, pt, w_len*pw] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [N*C*t_len, pt, h_len*ph, w_len*pw] + x = ggml_reshape_4d(ctx, x, pw * w_len, ph * h_len, pt * t_len, C * N); // [N*C, t_len*pt, h_len*ph, w_len*pw] return x; } @@ -707,7 +708,7 @@ namespace WAN { // patch_embedding x = patch_embedding->forward(ctx, x); // [N*dim, t_len, h_len, w_len] x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0] * x->ne[1] * x->ne[2], x->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] ggml_tensor* audio_local = nullptr; ggml_tensor* audio_global = nullptr; @@ -725,26 +726,26 @@ namespace WAN { // video tokens get cond_mask[0], reference tokens cond_mask[1] auto cond_mask = params["trainable_cond_mask.weight"]; - auto cm0 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 0, 1), config.dim, 1, 1); + auto cm0 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx, cond_mask, 1, 0, 1), config.dim, 1, 1); x = ggml_add(ctx->ggml_ctx, x, cm0); if (reference_latent != nullptr) { t_ref_len = reference_latent->ne[2]; auto ref = patch_embedding->forward(ctx, reference_latent); ref = ggml_reshape_3d(ctx->ggml_ctx, ref, ref->ne[0] * ref->ne[1] * ref->ne[2], ref->ne[3] / N, N); - ref = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, ref, 1, 0, 2, 3)); // [N, t_ref*h_len*w_len, dim] - auto cm1 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx->ggml_ctx, cond_mask, 1, 1, 2), config.dim, 1, 1); + ref = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, ref, 1, 0, 2, 3)); // [N, t_ref*h_len*w_len, dim] + auto cm1 = ggml_reshape_3d(ctx->ggml_ctx, ggml_ext_slice(ctx, cond_mask, 1, 1, 2), config.dim, 1, 1); ref = ggml_add(ctx->ggml_ctx, ref, cm1); x = ggml_concat(ctx->ggml_ctx, x, ref, 1); // Reference tokens use timestep 0. GGML_ASSERT(timestep->ne[0] == T); - timestep = ggml_ext_pad(ctx->ggml_ctx, timestep, (int)t_ref_len, 0, 0, 0); + timestep = ggml_ext_pad(ctx, timestep, (int)t_ref_len, 0, 0, 0); } } // time_embedding - auto e = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, config.freq_dim); + auto e = ggml_ext_timestep_embedding(ctx, timestep, config.freq_dim); e = time_embedding_0->forward(ctx, e); e = ggml_silu_inplace(ctx->ggml_ctx, e); e = time_embedding_2->forward(ctx, e); // [N, dim] or [N, T, dim] @@ -755,7 +756,7 @@ namespace WAN { e0 = ggml_reshape_4d(ctx->ggml_ctx, e0, e0->ne[0] / 6, 6, e0->ne[1], e0->ne[2]); // [N, 6, dim] or [N, T, 6, dim] context = text_embedding_0->forward(ctx, context); - context = ggml_ext_gelu(ctx->ggml_ctx, context); + context = ggml_ext_gelu(ctx, context); context = text_embedding_2->forward(ctx, context); // [N, context_txt_len, dim] int64_t context_img_len = 0; @@ -775,7 +776,7 @@ namespace WAN { c = vace_patch_embedding->forward(ctx, vace_context); // [N*dim, t_len, h_len, w_len] c = ggml_reshape_3d(ctx->ggml_ctx, c, c->ne[0] * c->ne[1] * c->ne[2], c->ne[3] / N, N); // [N, dim, t_len*h_len*w_len] - c = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] + c = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, c, 1, 0, 2, 3)); // [N, t_len*h_len*w_len, dim] } sd::ggml_graph_cut::mark_graph_cut(x, "wan.prelude", "x"); // sd::ggml_graph_cut::mark_graph_cut(e, "wan.prelude", "e"); @@ -806,7 +807,7 @@ namespace WAN { auto result = vace_block->forward(ctx, c, x_orig, e0, pe, context, context_img_len); auto c_skip = result.first; c = result.second; - c_skip = ggml_ext_scale(ctx->ggml_ctx, c_skip, vace_strength); + c_skip = ggml_ext_scale(ctx, c_skip, vace_strength); x = ggml_add(ctx->ggml_ctx, x, c_skip); } @@ -870,16 +871,16 @@ namespace WAN { if (reference_latent != nullptr) { // Exclude reference tokens from the generated video. - out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, t_len * h_len * w_len); + out = ggml_ext_slice(ctx, out, 1, 0, t_len * h_len * w_len); } - out = unpatchify(ctx->ggml_ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] + out = unpatchify(ctx, out, t_len, h_len, w_len); // [N*C, (T+pad_t) + (T2+pad_t2), H + pad_h, W + pad_w] // slice - out = ggml_ext_slice(ctx->ggml_ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] - out = ggml_ext_slice(ctx->ggml_ctx, out, 0, 0, W); // [N*C, T, H, W] + out = ggml_ext_slice(ctx, out, 2, 0, T); // [N*C, T, H + pad_h, W + pad_w] + out = ggml_ext_slice(ctx, out, 1, 0, H); // [N*C, T, H, W + pad_w] + out = ggml_ext_slice(ctx, out, 0, 0, W); // [N*C, T, H, W] return out; } diff --git a/src/model/diffusion/wan_audio.hpp b/src/model/diffusion/wan_audio.hpp index 28f31bda9f..387bb39e32 100644 --- a/src/model/diffusion/wan_audio.hpp +++ b/src/model/diffusion/wan_audio.hpp @@ -27,7 +27,7 @@ namespace WAN { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { // Replicate the first sample for causal left padding. if (kernel_size_ > 1) { - auto first = ggml_ext_slice(ctx->ggml_ctx, x, 0, 0, 1); + auto first = ggml_ext_slice(ctx, x, 0, 0, 1); for (int i = 0; i < kernel_size_ - 1; i++) { x = ggml_concat(ctx->ggml_ctx, first, x, 0); } @@ -57,7 +57,7 @@ namespace WAN { x = std::dynamic_pointer_cast(blocks[norm_key])->forward(ctx, x); x = ggml_silu(ctx->ggml_ctx, x); if (to_conv_layout) { - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); } return x; } @@ -87,11 +87,11 @@ namespace WAN { auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); std::vector tokens; // Each token group is normalized independently over channels. - for (auto& group : ggml_ext_chunk(ctx->ggml_ctx, local, num_token_, 1)) { + for (auto& group : ggml_ext_chunk(ctx, local, num_token_, 1)) { ggml_tensor* s = ggml_permute(ctx->ggml_ctx, group, 1, 0, 2, 3); s = norm1->forward(ctx, s); s = ggml_silu(ctx->ggml_ctx, s); - s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3)); + s = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, s, 1, 0, 2, 3)); s = conv_norm_silu(ctx, s, "conv2", "norm2", true); s = conv_norm_silu(ctx, s, "conv3", "norm3", false); tokens.push_back(ggml_reshape_3d(ctx->ggml_ctx, s, s->ne[0], 1, s->ne[1])); @@ -99,7 +99,7 @@ namespace WAN { auto padding = ggml_reshape_3d(ctx->ggml_ctx, params["padding_tokens"], hidden_dim_, 1, 1); padding = ggml_repeat(ctx->ggml_ctx, padding, tokens[0]); tokens.push_back(padding); - ggml_tensor* local_out = ggml_ext_vec_concat(ctx->ggml_ctx, tokens, 1); + ggml_tensor* local_out = ggml_ext_vec_concat(ctx, tokens, 1); if (!need_global_) { return {local_out, nullptr}; @@ -141,10 +141,10 @@ namespace WAN { auto x = ggml_mul(ctx->ggml_ctx, features, ggml_reshape_3d(ctx->ggml_ctx, weights, 1, 1, num_layers_)); x = ggml_div(ctx->ggml_ctx, x, ggml_sum(ctx->ggml_ctx, weights)); // Move the layer axis to ggml dimension 0 for reduction. - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); x = ggml_sum_rows(ctx->ggml_ctx, x); x = ggml_reshape_2d(ctx->ggml_ctx, x, x->ne[1], x->ne[2]); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); return std::dynamic_pointer_cast(blocks["encoder"])->forward(ctx, x); } }; @@ -187,12 +187,12 @@ namespace WAN { auto temb = ggml_silu(ctx->ggml_ctx, audio_global); temb = adain_linear->forward(ctx, temb); - auto shift = ggml_ext_slice(ctx->ggml_ctx, temb, 0, 0, dim_); - auto scale = ggml_ext_slice(ctx->ggml_ctx, temb, 0, dim_, dim_ * 2); + auto shift = ggml_ext_slice(ctx, temb, 0, 0, dim_); + auto scale = ggml_ext_slice(ctx, temb, 0, dim_, dim_ * 2); shift = ggml_reshape_3d(ctx->ggml_ctx, shift, dim_, 1, T); scale = ggml_reshape_3d(ctx->ggml_ctx, scale, dim_, 1, T); - auto x_vid = ggml_ext_slice(ctx->ggml_ctx, x, 1, 0, seq_len); + auto x_vid = ggml_ext_slice(ctx, x, 1, 0, seq_len); auto h = ggml_reshape_3d(ctx->ggml_ctx, x_vid, dim_, n_tok, T); h = adain_norm->forward(ctx, h); h = ggml_add(ctx->ggml_ctx, h, ggml_mul(ctx->ggml_ctx, h, scale)); @@ -203,7 +203,7 @@ namespace WAN { auto x_head = ggml_add(ctx->ggml_ctx, x_vid, res); if (seq_len < n_token) { - auto x_tail = ggml_ext_slice(ctx->ggml_ctx, x, 1, seq_len, n_token); + auto x_tail = ggml_ext_slice(ctx, x, 1, seq_len, n_token); return ggml_concat(ctx->ggml_ctx, x_head, x_tail, 1); } return x_head; diff --git a/src/model/diffusion/z_image.hpp b/src/model/diffusion/z_image.hpp index bc9ddc4498..43df4d342a 100644 --- a/src/model/diffusion/z_image.hpp +++ b/src/model/diffusion/z_image.hpp @@ -179,12 +179,10 @@ namespace ZImage { auto k_proj = std::dynamic_pointer_cast(blocks["to_k"]); auto v_proj = std::dynamic_pointer_cast(blocks["to_v"]); - if (sd_backend_is(ctx->backend, "ROCm")) { - out_proj->set_force_prec_f32(true); - q_proj->set_force_prec_f32(true); - k_proj->set_force_prec_f32(true); - v_proj->set_force_prec_f32(true); - } + out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); + q_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); + k_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); + v_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); auto q = ggml_reshape_4d(ctx->ggml_ctx, q_proj->forward(ctx, x), head_dim, num_heads, n_token, N); auto k = ggml_reshape_4d(ctx->ggml_ctx, k_proj->forward(ctx, x), head_dim, num_kv_heads, n_token, N); @@ -201,10 +199,8 @@ namespace ZImage { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); - if (sd_backend_is(ctx->backend, "ROCm")) { - out_proj->set_force_prec_f32(true); - qkv_proj->set_force_prec_f32(true); - } + out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); + qkv_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim] qkv = ggml_reshape_4d(ctx->ggml_ctx, qkv, head_dim, num_heads + num_kv_heads * 2, qkv->ne[1], qkv->ne[2]); // [N, n_token, num_heads + num_kv_heads*2, head_dim] @@ -281,9 +277,7 @@ namespace ZImage { auto w2 = std::dynamic_pointer_cast(blocks["w2"]); auto w3 = std::dynamic_pointer_cast(blocks["w3"]); - if (sd_backend_is(ctx->backend, "Vulkan") || sd_backend_is(ctx->backend, "ROCm")) { - w2->set_force_prec_f32(true); - } + w2->set_precision_flags(GGML_PRECISION_F32_VULKAN | GGML_PRECISION_F32_ROCM); auto x1 = w1->forward(ctx, x); auto x3 = w3->forward(ctx, x); @@ -353,7 +347,7 @@ namespace ZImage { auto adaLN_modulation_0 = std::dynamic_pointer_cast(blocks["adaLN_modulation.0"]); auto m = adaLN_modulation_0->forward(ctx, adaln_input); // [N, 4 * hidden_size] - auto mods = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0); + auto mods = ggml_ext_chunk(ctx, m, 4, 0); auto scale_msa = mods[0]; auto gate_msa = mods[1]; auto scale_mlp = mods[2]; @@ -457,12 +451,10 @@ namespace ZImage { 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); - } + out_proj->set_precision_flags(GGML_PRECISION_F32_ROCM); + to_q->set_precision_flags(GGML_PRECISION_F32_ROCM); + to_k->set_precision_flags(GGML_PRECISION_F32_ROCM); + to_v->set_precision_flags(GGML_PRECISION_F32_ROCM); 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); @@ -497,7 +489,7 @@ namespace ZImage { } auto m = adaLN_modulation_0->forward(ctx, t_emb); - auto mods = ggml_ext_chunk(ctx->ggml_ctx, m, 4, 0); + auto mods = ggml_ext_chunk(ctx, m, 4, 0); auto scale_msa = mods[0]; auto gate_msa = mods[1]; auto scale_mlp = mods[2]; @@ -713,7 +705,7 @@ namespace ZImage { 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)); + ref_hints.push_back(ggml_ext_scale(ctx, out.first, strength)); sd::ggml_graph_cut::mark_graph_cut(c, "z_image.control.noise_refiner." + std::to_string(i), "control"); } @@ -757,7 +749,7 @@ namespace ZImage { 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)); + main_hints.push_back(ggml_ext_scale(ctx, out.first, strength)); sd::ggml_graph_cut::mark_graph_cut(ctrl_unified, "z_image.control.layers." + std::to_string(i), "control"); } @@ -807,8 +799,8 @@ namespace ZImage { GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]); - 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]); + auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]); + auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]); std::vector ref_hints; ggml_tensor* control_ctx = nullptr; @@ -858,7 +850,7 @@ namespace ZImage { txt_img = final_layer->forward(ctx, txt_img, t_emb); // [N, n_txt_token + n_txt_pad_token + n_img_token + n_img_pad_token, ph*pw*C] - img = ggml_ext_slice(ctx->ggml_ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); // [N, n_img_token, ph*pw*C] + img = ggml_ext_slice(ctx, txt_img, 1, n_txt_token + n_txt_pad_token, n_txt_token + n_txt_pad_token + n_img_token); // [N, n_img_token, ph*pw*C] return img; } @@ -906,10 +898,10 @@ namespace ZImage { 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] + out = ggml_ext_slice(ctx, out, 1, 0, n_img_token); // [N, n_img_token, ph*pw*C] + out = DiT::unpatchify_and_crop(ctx, out, H, W, patch_size, patch_size, false); // [N, C, H, W] - out = ggml_ext_scale(ctx->ggml_ctx, out, -1.f); + out = ggml_ext_scale(ctx, out, -1.f); return out; } diff --git a/src/model/diffusion/z_image_l2p.hpp b/src/model/diffusion/z_image_l2p.hpp index 2ba6680329..7146295a46 100644 --- a/src/model/diffusion/z_image_l2p.hpp +++ b/src/model/diffusion/z_image_l2p.hpp @@ -134,7 +134,7 @@ namespace ZImageL2P { const int64_t n_img = w * h; const int64_t n_txt = context->ne[1]; - auto img = DiT::patchify(gctx, x, config.patch_size, config.patch_size, false); + auto img = DiT::patchify(ctx, x, config.patch_size, config.patch_size, false); img = std::dynamic_pointer_cast(blocks["x_embedder"])->forward(ctx, img); auto txt = std::dynamic_pointer_cast(blocks["cap_embedder.0"])->forward(ctx, context); txt = std::dynamic_pointer_cast(blocks["cap_embedder.1"])->forward(ctx, txt); @@ -155,8 +155,8 @@ namespace ZImageL2P { } GGML_ASSERT(txt->ne[1] + img->ne[1] == pe->ne[3]); - auto txt_pe = ggml_ext_slice(gctx, pe, 3, 0, txt->ne[1]); - auto img_pe = ggml_ext_slice(gctx, pe, 3, txt->ne[1], pe->ne[3]); + auto txt_pe = ggml_ext_slice(ctx, pe, 3, 0, txt->ne[1]); + auto img_pe = ggml_ext_slice(ctx, pe, 3, txt->ne[1], pe->ne[3]); for (int64_t i = 0; i < config.num_refiner_layers; ++i) { txt = std::dynamic_pointer_cast(blocks["context_refiner." + std::to_string(i)])->forward(ctx, txt, txt_pe); sd::ggml_graph_cut::mark_graph_cut(txt, "z_image_l2p.context_refiner." + std::to_string(i), "txt"); @@ -175,11 +175,11 @@ namespace ZImageL2P { // The local decoder consumes the raw last-layer hidden states: there is no final // norm or adaLN modulation in front of it, unlike Z-Image's final_layer. const int64_t img_start = n_txt + n_txt_pad; - auto feat = ggml_ext_slice(gctx, txt_img, 1, img_start, img_start + n_img); // [N, h*w, hidden_size] + auto feat = ggml_ext_slice(ctx, txt_img, 1, img_start, img_start + n_img); // [N, h*w, hidden_size] feat = ggml_reshape_4d(gctx, feat, config.hidden_size, w, h, N); // [N, h, w, hidden_size] feat = ggml_cont(gctx, ggml_permute(gctx, feat, 2, 0, 1, 3)); // [N, hidden_size, h, w] auto out = std::dynamic_pointer_cast(blocks["local_decoder"])->forward(ctx, x, feat); // [N, C, H, W] - return ggml_ext_scale(gctx, out, -1.f); + return ggml_ext_scale(ctx, out, -1.f); } }; diff --git a/src/model/te/clip.hpp b/src/model/te/clip.hpp index cac613eda0..858895d3c7 100644 --- a/src/model/te/clip.hpp +++ b/src/model/te/clip.hpp @@ -35,9 +35,9 @@ struct CLIPMLP : public GGMLBlock { x = fc1->forward(ctx, x); if (use_gelu) { - x = ggml_ext_gelu(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu(ctx, x, true); } else { - x = ggml_ext_gelu_quick(ctx->ggml_ctx, x, true); + x = ggml_ext_gelu_quick(ctx, x, true); } x = fc2->forward(ctx, x); return x; @@ -228,10 +228,10 @@ class CLIPVisionEmbeddings : public GGMLBlock { // concat(patch_embedding, class_embedding) + position_embedding ggml_tensor* patch_embedding; int64_t N = pixel_values->ne[3]; - patch_embedding = ggml_ext_conv_2d(ctx->ggml_ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size); // [N, embed_dim, image_size // pacht_size, image_size // pacht_size] - patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N); // [N, embed_dim, num_patches] - patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3)); // [N, num_patches, embed_dim] - patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N); // [N, num_patches, embed_dim, 1] + patch_embedding = ggml_ext_conv_2d(ctx, pixel_values, patch_embed_weight, nullptr, patch_size, patch_size); // [N, embed_dim, image_size // pacht_size, image_size // pacht_size] + patch_embedding = ggml_reshape_3d(ctx->ggml_ctx, patch_embedding, num_patches, embed_dim, N); // [N, embed_dim, num_patches] + patch_embedding = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, patch_embedding, 1, 0, 2, 3)); // [N, num_patches, embed_dim] + patch_embedding = ggml_reshape_4d(ctx->ggml_ctx, patch_embedding, 1, embed_dim, num_patches, N); // [N, num_patches, embed_dim, 1] ggml_tensor* class_embedding = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, embed_dim, N); class_embedding = ggml_repeat(ctx->ggml_ctx, class_embed_weight, class_embedding); // [N, embed_dim] @@ -325,7 +325,7 @@ class CLIPTextModel : public GGMLBlock { auto text_projection = params["text_projection"]; ggml_tensor* pooled = ggml_view_1d(ctx->ggml_ctx, x, hidden_size, x->nb[1] * max_token_idx); if (text_projection != nullptr) { - pooled = ggml_ext_linear(ctx->ggml_ctx, pooled, text_projection, nullptr); + pooled = ggml_ext_linear(ctx, pooled, text_projection, nullptr); } else { LOG_VERBOSE("identity projection"); } @@ -426,7 +426,7 @@ class CLIPProjection : public UnaryBlock { if (transpose_weight) { w = ggml_cont(ctx->ggml_ctx, ggml_transpose(ctx->ggml_ctx, w)); } - return ggml_ext_linear(ctx->ggml_ctx, x, w, nullptr); + return ggml_ext_linear(ctx, x, w, nullptr); } }; diff --git a/src/model/te/llada_image_te.hpp b/src/model/te/llada_image_te.hpp index 2da78413b1..64c18567af 100644 --- a/src/model/te/llada_image_te.hpp +++ b/src/model/te/llada_image_te.hpp @@ -78,15 +78,15 @@ namespace LLaDAImageTE { auto b = params["in_proj_bias"]; auto slice_w = [&](int64_t index) { - return ggml_ext_slice(gctx, w, 1, index * hidden_size, (index + 1) * hidden_size); + return ggml_ext_slice(ctx, w, 1, index * hidden_size, (index + 1) * hidden_size); }; auto slice_b = [&](int64_t index) { - return ggml_ext_slice(gctx, b, 0, index * hidden_size, (index + 1) * hidden_size); + return ggml_ext_slice(ctx, b, 0, index * hidden_size, (index + 1) * hidden_size); }; - auto q = ggml_ext_linear(gctx, query, slice_w(0), slice_b(0)); - auto k = ggml_ext_linear(gctx, context, slice_w(1), slice_b(1)); - auto v = ggml_ext_linear(gctx, context, slice_w(2), slice_b(2)); + auto q = ggml_ext_linear(ctx, query, slice_w(0), slice_b(0)); + auto k = ggml_ext_linear(ctx, context, slice_w(1), slice_b(1)); + auto v = ggml_ext_linear(ctx, context, slice_w(2), slice_b(2)); auto x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, num_queries, hidden_size] return out_proj->forward(ctx, x); @@ -297,9 +297,9 @@ namespace LLaDAImageTE { int64_t hidden_size = num_heads * head_dim; auto qkv = qkv_proj->forward(ctx, x); - auto q = ggml_ext_slice(gctx, qkv, 0, 0, hidden_size); - auto k = ggml_ext_slice(gctx, qkv, 0, hidden_size, hidden_size * 2); - auto v = ggml_ext_slice(gctx, qkv, 0, hidden_size * 2, hidden_size * 3); + auto q = ggml_ext_slice(ctx, qkv, 0, 0, hidden_size); + auto k = ggml_ext_slice(ctx, qkv, 0, hidden_size, hidden_size * 2); + auto v = ggml_ext_slice(ctx, qkv, 0, hidden_size * 2, hidden_size * 3); auto out = ggml_ext_attention_ext(ctx, q, k, v, num_heads); return out_proj->forward(ctx, out); diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index add8b6ff9a..7ce1f175e3 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -485,7 +485,7 @@ namespace LLM { 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"); + w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight"); } x = ggml_rms_norm(ctx->ggml_ctx, x, eps); auto scaled = ggml_mul(ctx->ggml_ctx, x, w); @@ -941,7 +941,7 @@ namespace LLM { // 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. if (txt_token_end > txt_token_start) { - auto txt_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, txt_token_start, txt_token_end); + auto txt_embed = ggml_ext_slice(ctx, raw_x, 1, txt_token_start, txt_token_end); if (input_embed == nullptr) { input_embed = txt_embed; } else { @@ -960,7 +960,7 @@ namespace LLM { 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); + auto final_txt_embed = ggml_ext_slice(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] + removed_tokens == input_embed->ne[1]); @@ -982,16 +982,16 @@ namespace LLM { GGML_ASSERT(index >= token_start); GGML_ASSERT(index + image_embed->ne[1] <= raw_x->ne[1]); if (index > token_start) { - auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, index); + auto text_embed = ggml_ext_slice(ctx, raw_x, 1, token_start, index); output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1); } - auto visual_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, index, index + image_embed->ne[1]); + auto visual_embed = ggml_ext_slice(ctx, raw_x, 1, index, index + image_embed->ne[1]); visual_embed = ggml_add(ctx->ggml_ctx, visual_embed, image_embed); output = output == nullptr ? visual_embed : ggml_concat(ctx->ggml_ctx, output, visual_embed, 1); token_start = index + image_embed->ne[1]; } if (token_start < raw_x->ne[1]) { - auto text_embed = ggml_ext_slice(ctx->ggml_ctx, raw_x, 1, token_start, raw_x->ne[1]); + auto text_embed = ggml_ext_slice(ctx, raw_x, 1, token_start, raw_x->ne[1]); output = output == nullptr ? text_embed : ggml_concat(ctx->ggml_ctx, output, text_embed, 1); } GGML_ASSERT(output != nullptr && output->ne[1] == raw_x->ne[1]); @@ -1020,7 +1020,7 @@ namespace LLM { auto linear_fc1 = std::dynamic_pointer_cast(blocks["linear_fc1"]); auto linear_fc2 = std::dynamic_pointer_cast(blocks["linear_fc2"]); x = linear_fc1->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = ggml_ext_gelu(ctx, x); x = linear_fc2->forward(ctx, x); } else { auto gate_proj = std::dynamic_pointer_cast(blocks["gate_proj"]); @@ -1105,11 +1105,11 @@ namespace LLM { auto proj_0 = std::dynamic_pointer_cast(blocks["proj.0"]); auto proj_1 = std::dynamic_pointer_cast(blocks["proj.1"]); - auto x0 = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto x0 = ggml_ext_slice(ctx, x, 2, 0, 1); x0 = ggml_reshape_4d(ctx->ggml_ctx, x0, x0->ne[0], x0->ne[1], in_channels, x0->ne[3] / in_channels); x0 = proj_0->forward(ctx, x0); - auto x1 = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, 2); + auto x1 = ggml_ext_slice(ctx, x, 2, 1, 2); x1 = ggml_reshape_4d(ctx->ggml_ctx, x1, x1->ne[0], x1->ne[1], in_channels, x1->ne[3] / in_channels); x1 = proj_1->forward(ctx, x1); @@ -1173,7 +1173,7 @@ namespace LLM { x = ln_q->forward(ctx, x); x = ggml_reshape_2d(ctx->ggml_ctx, x, hidden_size, ggml_nelements(x) / hidden_size); x = mlp_0->forward(ctx, x); - x = ggml_ext_gelu(ctx->ggml_ctx, x); + x = ggml_ext_gelu(ctx, x); x = mlp_2->forward(ctx, x); return x; } @@ -1252,7 +1252,7 @@ namespace LLM { } else { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); auto qkv = qkv_proj->forward(ctx, x); - qkv_vec = split_qkv(ctx->ggml_ctx, qkv); + qkv_vec = ggml_ext_split_qkv(ctx, qkv); } auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); // [N, n_token, n_head, d_head] @@ -1487,12 +1487,12 @@ namespace LLM { for (int i = 0; i < kernel; ++i) { auto w = ggml_cont(c, ggml_view_2d(c, params["conv1d.weight"], 1, conv_dim, kernel * sizeof(float), i * sizeof(float))); w = ggml_reshape_1d(c, w, conv_dim); - auto part = ggml_mul(c, ggml_ext_slice(c, padded, 1, i, i + n), w); + auto part = ggml_mul(c, ggml_ext_slice(ctx, padded, 1, i, i + n), w); mixed = mixed ? ggml_add(c, mixed, part) : part; } mixed = ggml_silu(c, mixed); auto heads = [&](int start, int count) { - auto h = ggml_ext_slice(c, mixed, 0, start, start + head_dim * count); + auto h = ggml_ext_slice(ctx, mixed, 0, start, start + head_dim * count); return ggml_reshape_4d(c, h, head_dim, count, n, 1); }; auto repeat_keys = [&](ggml_tensor* h) { @@ -1514,7 +1514,7 @@ namespace LLM { auto a = gguf_layout ? params["A"] : ggml_neg(c, ggml_exp(c, params["A_log"])); decay = ggml_mul(c, decay, a); decay = ggml_reshape_4d(c, decay, 1, value_heads, n, 1); - auto state = ggml_ext_zeros(c, head_dim, head_dim, value_heads, 1); + auto state = ggml_ext_zeros(ctx, head_dim, head_dim, value_heads, 1); auto out = ggml_gated_delta_net(c, q, k, v, decay, beta, state, 1); 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); @@ -1594,13 +1594,14 @@ namespace LLM { // ggml applies theta_i / freq_factors[i], so a huge factor collapses the angle to zero and // leaves that pair unrotated. This reproduces transformers' "proportional" RoPE, whose // inv_freq is zero-padded past `rope_pairs`, without reordering the head. - ggml_tensor* rope_freq_factors(ggml_context* ctx) const { + ggml_tensor* rope_freq_factors(GGMLRunnerContext* runner_ctx) const { + auto ctx = runner_ctx->ggml_ctx; int pairs = head_dim / 2; if (rope_pairs >= pairs) { return nullptr; } - auto rotated = ggml_ext_ones(ctx, rope_pairs, 1, 1, 1); - auto unrotated = ggml_ext_full(ctx, 1e30f, pairs - rope_pairs, 1, 1, 1); + auto rotated = ggml_ext_ones(runner_ctx, rope_pairs, 1, 1, 1); + auto unrotated = ggml_ext_full(runner_ctx, 1e30f, pairs - rope_pairs, 1, 1, 1); return ggml_concat(ctx, rotated, unrotated, 0); } @@ -1624,9 +1625,9 @@ namespace LLM { auto qkv = qkv_proj->forward(ctx, x); // [N, n_token, (num_heads + num_kv_heads*2)*head_dim] int64_t q_len = num_heads * head_dim; int64_t k_len = num_kv_heads * head_dim; - q = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, 0, q_len); - k = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len, q_len + k_len); - v = ggml_ext_slice(ctx->ggml_ctx, qkv, 0, q_len + k_len, q_len + k_len * 2); + q = ggml_ext_slice(ctx, qkv, 0, 0, q_len); + k = ggml_ext_slice(ctx, qkv, 0, q_len, q_len + k_len); + v = ggml_ext_slice(ctx, qkv, 0, q_len + k_len, q_len + k_len * 2); } else { auto q_proj = std::dynamic_pointer_cast(blocks["q_proj"]); auto k_proj = std::dynamic_pointer_cast(blocks["k_proj"]); @@ -1640,9 +1641,9 @@ namespace LLM { ggml_tensor* output_gate = nullptr; if (arch == LLMArch::QWEN3_5) { q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim * 2, num_heads, n_token, N); - output_gate = ggml_ext_slice(ctx->ggml_ctx, q, 0, head_dim, head_dim * 2); + output_gate = ggml_ext_slice(ctx, q, 0, head_dim, head_dim * 2); output_gate = ggml_reshape_3d(ctx->ggml_ctx, output_gate, head_dim * num_heads, n_token, N); - q = ggml_ext_slice(ctx->ggml_ctx, q, 0, 0, head_dim); + q = ggml_ext_slice(ctx, q, 0, 0, head_dim); } q = ggml_reshape_4d(ctx->ggml_ctx, q, head_dim, num_heads, n_token, N); // [N, n_token, num_heads, head_dim] k = ggml_reshape_4d(ctx->ggml_ctx, k, head_dim, num_kv_heads, n_token, N); // [N, n_token, num_kv_heads, head_dim] @@ -1666,12 +1667,12 @@ namespace LLM { auto rotate = [&](ggml_tensor* input) { auto gctx = ctx->ggml_ctx; int64_t half = rope_cos->ne[0]; - auto first = ggml_ext_slice(gctx, input, 0, 0, half); - auto second = ggml_ext_slice(gctx, input, 0, half, 2 * half); + auto first = ggml_ext_slice(ctx, input, 0, 0, half); + auto second = ggml_ext_slice(ctx, input, 0, half, 2 * half); auto left = ggml_sub(gctx, ggml_mul(gctx, first, rope_cos), ggml_mul(gctx, second, rope_sin)); auto right = ggml_add(gctx, ggml_mul(gctx, second, rope_cos), ggml_mul(gctx, first, rope_sin)); auto rotated = ggml_concat(gctx, left, right, 0); - return ggml_concat(gctx, rotated, ggml_ext_slice(gctx, input, 0, 2 * half, input->ne[0]), 0); + return ggml_concat(gctx, rotated, ggml_ext_slice(ctx, input, 0, 2 * half, input->ne[0]), 0); }; q = rotate(q); k = rotate(k); @@ -1746,7 +1747,7 @@ namespace LLM { 1.f); } else if (arch == LLMArch::GEMMA4_12B) { float rope_theta = (rope_index == 1 ? 10000.0f : 1000000.0f); - auto freq_factors = rope_freq_factors(ctx->ggml_ctx); + auto freq_factors = rope_freq_factors(ctx); q = ggml_rope_ext(ctx->ggml_ctx, q, input_pos, @@ -1844,18 +1845,18 @@ namespace LLM { if (unscaled_attention) { // Gemma 4 attends with scaling=1.0; undo the helper's own 1/sqrt(head_dim). - q = ggml_ext_scale(ctx->ggml_ctx, q, std::sqrt(static_cast(head_dim))); + q = ggml_ext_scale(ctx, q, std::sqrt(static_cast(head_dim))); } - q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)); // [N, num_heads, n_token, head_dim] - q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]); // [N*num_heads, n_token, head_dim] + q = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, q, 0, 2, 1, 3)); // [N, num_heads, n_token, head_dim] + q = ggml_reshape_3d(ctx->ggml_ctx, q, q->ne[0], q->ne[1], q->ne[2] * q->ne[3]); // [N*num_heads, n_token, head_dim] - k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)); // [N, num_kv_heads, n_token, head_dim] - k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]); // [N*num_kv_heads, n_token, head_dim] + k = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, k, 0, 2, 1, 3)); // [N, num_kv_heads, n_token, head_dim] + k = ggml_reshape_3d(ctx->ggml_ctx, k, k->ne[0], k->ne[1], k->ne[2] * k->ne[3]); // [N*num_kv_heads, n_token, head_dim] if (arch == LLMArch::GPT_OSS_20B) { GGML_ASSERT(N == 1); - auto v_attn = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3)); // [N, kv_heads, head_dim, tokens] + auto v_attn = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, v, 1, 2, 0, 3)); // [N, kv_heads, head_dim, tokens] v_attn = ggml_reshape_3d(ctx->ggml_ctx, v_attn, n_token, head_dim, num_kv_heads * N); auto kq = ggml_mul_mat(ctx->ggml_ctx, k, q); @@ -1870,7 +1871,7 @@ namespace LLM { auto kqv = ggml_mul_mat(ctx->ggml_ctx, v_attn, kq); kqv = ggml_reshape_4d(ctx->ggml_ctx, kqv, head_dim, n_token, num_heads, N); kqv = ggml_permute(ctx->ggml_ctx, kqv, 0, 2, 1, 3); - x = ggml_ext_cont(ctx->ggml_ctx, kqv); + x = ggml_ext_cont(ctx, kqv); x = ggml_reshape_3d(ctx->ggml_ctx, x, head_dim * num_heads, n_token, N); } else { x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, attention_mask, true, ctx->flash_attn_enabled); // [N, n_token, hidden_size] @@ -2059,7 +2060,7 @@ namespace LLM { std::vector intermediate_outputs; if (config.normalize_input) { - x = ggml_ext_scale(ctx->ggml_ctx, x, std::sqrt(static_cast(config.hidden_size)), true); + x = ggml_ext_scale(ctx, x, std::sqrt(static_cast(config.hidden_size)), true); } if (return_all_hidden_states) { intermediate_outputs.push_back(x); @@ -2207,9 +2208,10 @@ namespace LLM { std::array, 4> pos_embed_idx_data_; std::array, 4> pos_embed_weight_data_; - static ggml_tensor* process_image_common(ggml_context* ctx, + static ggml_tensor* process_image_common(GGMLRunnerContext* runner_ctx, ggml_tensor* image, const LLMVisionConfig& vision_params) { + auto ctx = runner_ctx->ggml_ctx; // image: [C, H, W] // return: [grid_t*(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw], grid_t == 1 int64_t C = image->ne[2]; @@ -2222,16 +2224,16 @@ namespace LLM { int64_t pw = vision_params.patch_size; image = ggml_reshape_4d(ctx, image, pw, mw, (W / mw / pw), H * C); // [C*H, (W/mw/pw), mw, pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); // [mw, C*H, (W/mw/pw), pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 3, 1)); // [mw, C*H, (W/mw/pw), pw] image = ggml_reshape_4d(ctx, image, pw * (W / mw / pw), H, C, mw); // [mw, C, H, (W/mw/pw)*pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 3, 1)); // [H, mw, C, (W/mw/pw)*pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 3, 1)); // [H, mw, C, (W/mw/pw)*pw] image = ggml_reshape_4d(ctx, image, pw, (W / mw / pw) * C * mw, ph, mh * (H / mh / ph)); // [(H/mh/ph)*mh, ph, mw*C*(W/mw/pw), pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh, mw*C*(W/mw/pw), ph, pw] image = ggml_reshape_4d(ctx, image, pw * ph, (W / mw / pw), C, mw * mh * (H / mh / ph)); // [(H/mh/ph)*mh*mw, C, (W/mw/pw), ph*pw] image = ggml_concat(ctx, image, image, 0); // [(H/mh/ph)*mh*mw, C, (W/mw/pw), pt*ph*pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3)); // [(H/mh/ph)*mh*mw, (W/mw/pw), C, pt*ph*pw] image = ggml_reshape_4d(ctx, image, pw * ph * pt * C, (W / mw / pw), mw * mh, (H / mh / ph)); // [(H/mh/ph), mh*mw, (W/mw/pw), C*pt*ph*pw] - image = ggml_cont(ctx, ggml_ext_torch_permute(ctx, image, 0, 2, 1, 3)); // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw] + image = ggml_cont(ctx, ggml_ext_torch_permute(runner_ctx, image, 0, 2, 1, 3)); // [(H/mh/ph), (W/mw/pw), mh*mw, C*pt*ph*pw] image = ggml_reshape_2d(ctx, image, pw * ph * pt * C, mw * mh * (W / mw / pw) * (H / mh / ph)); // [(H/mh/ph)*(W/mw/pw)*mh*mw, C*pt*ph*pw] return image; } @@ -2301,7 +2303,6 @@ namespace LLM { } static ggml_tensor* encode_image_common(GGMLRunner* runner, - ggml_context* compute_ctx, GGMLRunnerContext* runner_ctx, ggml_tensor* image, const LLMVisionConfig& vision_params, @@ -2313,13 +2314,14 @@ namespace LLM { std::array, 4>& pos_embed_idx_data, std::array, 4>& pos_embed_weight_data, std::vector* output_tensors = nullptr) { + auto compute_ctx = runner_ctx->ggml_ctx; GGML_ASSERT(image->ne[1] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0); GGML_ASSERT(image->ne[0] % (vision_params.patch_size * vision_params.spatial_merge_size) == 0); int grid_h = static_cast(image->ne[1]) / vision_params.patch_size; int grid_w = static_cast(image->ne[0]) / vision_params.patch_size; - auto pixel_values = process_image_common(compute_ctx, image, vision_params); + auto pixel_values = process_image_common(runner_ctx, image, vision_params); int head_dim = static_cast(vision_params.hidden_size / vision_params.num_heads); if (vision_params.arch == LLMVisionArch::QWEN3_VL) { @@ -2678,8 +2680,8 @@ namespace LLM { return grid_t * grid_h * grid_w; } - ggml_tensor* process_image(ggml_context* ctx, ggml_tensor* image) { - return process_image_common(ctx, image, config.vision); + ggml_tensor* process_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { + return process_image_common(runner_ctx, image, config.vision); } ggml_tensor* build_patch_pos_embeds(GGMLRunnerContext* runner_ctx, @@ -2698,7 +2700,6 @@ namespace LLM { ggml_tensor* encode_image(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { return encode_image_common(this, - compute_ctx, runner_ctx, image, config.vision, @@ -2714,7 +2715,6 @@ namespace LLM { std::vector encode_image_outputs(GGMLRunnerContext* runner_ctx, ggml_tensor* image) { std::vector outputs; encode_image_common(this, - compute_ctx, runner_ctx, image, config.vision, diff --git a/src/model/te/ming_image_te.hpp b/src/model/te/ming_image_te.hpp index ab3b8e6dbc..c8f20476f8 100644 --- a/src/model/te/ming_image_te.hpp +++ b/src/model/te/ming_image_te.hpp @@ -73,16 +73,16 @@ namespace MingImageTE { auto connector = std::dynamic_pointer_cast(blocks["connector"]); auto x = backbone->embed(ctx, ids); const int64_t query_start = prompt_length + 1; - auto before = ggml_ext_slice(gctx, x, 1, 0, query_start); - auto after = ggml_ext_slice(gctx, x, 1, query_start + config.num_queries, x->ne[1]); + auto before = ggml_ext_slice(ctx, x, 1, 0, query_start); + auto after = ggml_ext_slice(ctx, x, 1, query_start + config.num_queries, x->ne[1]); x = ggml_concat(gctx, ggml_concat(gctx, before, params["query_tokens_dict.16x16"], 1), after, 1); // HF hidden_states[20] includes final RMSNorm; sd.cpp selects it as num_layers + 1. x = backbone->forward_embeds(ctx, x, nullptr, mask, {5, 12, 21}, {}, nullptr, false, image_mask, cos, sin); - auto direct = ggml_ext_slice(gctx, x, 1, 0, prompt_length); + auto direct = ggml_ext_slice(ctx, x, 1, 0, prompt_length); direct = std::dynamic_pointer_cast(blocks["proj_directvlm.0"])->forward(ctx, direct); direct = std::dynamic_pointer_cast(blocks["proj_directvlm.1"])->forward(ctx, direct); - auto queries = ggml_ext_slice(gctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3); - queries = ggml_ext_slice(gctx, queries, 1, query_start, query_start + config.num_queries); + auto queries = ggml_ext_slice(ctx, x, 0, config.backbone.hidden_size * 2, config.backbone.hidden_size * 3); + queries = ggml_ext_slice(ctx, queries, 1, query_start, query_start + config.num_queries); queries = std::dynamic_pointer_cast(blocks["proj_in"])->forward(ctx, queries); queries = connector->forward_embeds(ctx, queries, connector_positions, nullptr, {}); queries = std::dynamic_pointer_cast(blocks["proj_out"])->forward(ctx, queries); diff --git a/src/model/te/ph3_image_connector.hpp b/src/model/te/ph3_image_connector.hpp index 1e74e8e765..a0e82eca2a 100644 --- a/src/model/te/ph3_image_connector.hpp +++ b/src/model/te/ph3_image_connector.hpp @@ -28,11 +28,11 @@ struct PH3ImageTextConnector : public UnaryBlock { 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)); + auto sink = ggml_repeat(ctx->ggml_ctx, params["sink_embedding"], ggml_ext_slice(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); + return ggml_concat(ctx->ggml_ctx, sink, ggml_ext_slice(ctx, x, 1, 1, x->ne[1]), 1); } }; diff --git a/src/model/te/t5.hpp b/src/model/te/t5.hpp index 29c6f87311..c8e9b47afc 100644 --- a/src/model/te/t5.hpp +++ b/src/model/te/t5.hpp @@ -249,7 +249,7 @@ class T5Attention : public GGMLBlock { } } - k = ggml_ext_scale(ctx->ggml_ctx, k, ::sqrtf(static_cast(d_head)), true); + k = ggml_ext_scale(ctx, k, ::sqrtf(static_cast(d_head)), true); x = ggml_ext_attention_ext(ctx, q, k, v, num_heads, mask); // [N, n_token, d_head * n_head] diff --git a/src/model/upscaler/esrgan.hpp b/src/model/upscaler/esrgan.hpp index 66029f5deb..d132f4fc79 100644 --- a/src/model/upscaler/esrgan.hpp +++ b/src/model/upscaler/esrgan.hpp @@ -127,7 +127,7 @@ class ResidualDenseBlock : public GGMLBlock { x_cat = ggml_concat(ctx->ggml_ctx, x_cat, x4, 2); auto x5 = conv5->forward(ctx, x_cat); - x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, x5, 0.2f), x); + x5 = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx, x5, 0.2f), x); return x5; } }; @@ -152,7 +152,7 @@ class RRDB : public GGMLBlock { out = rdb2->forward(ctx, out); out = rdb3->forward(ctx, out); - out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx->ggml_ctx, out, 0.2f), x); + out = ggml_add(ctx->ggml_ctx, ggml_ext_scale(ctx, out, 0.2f), x); return out; } }; diff --git a/src/model/upscaler/flash_vsr.hpp b/src/model/upscaler/flash_vsr.hpp index 470bfc3065..addfc532ad 100644 --- a/src/model/upscaler/flash_vsr.hpp +++ b/src/model/upscaler/flash_vsr.hpp @@ -146,7 +146,7 @@ namespace RelayVSR { auto y = attention->forward(ctx, attention_norm->forward(ctx, x), qpe, kpe, local, name, condition_slot, start_slot, end_slot, history); x = ggml_add(c, x, y); y = feed_forward_gate_up_proj->forward(ctx, feed_forward_norm->forward(ctx, x)); - auto parts = ggml_ext_chunk(c, y, 2, 0); + auto parts = ggml_ext_chunk(ctx, y, 2, 0); y = ggml_mul(c, ggml_silu(c, parts[0]), parts[1]); return ggml_add(c, x, feed_forward_down->forward(ctx, y)); } diff --git a/src/model/upscaler/h3_latent_upscaler.hpp b/src/model/upscaler/h3_latent_upscaler.hpp index 80b96e7395..0411ca1268 100644 --- a/src/model/upscaler/h3_latent_upscaler.hpp +++ b/src/model/upscaler/h3_latent_upscaler.hpp @@ -195,8 +195,8 @@ namespace H3VUpsampler { 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* scale = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(ctx, emb_1d, 0, 0, channels), 1, 1, 1, channels); + ggml_tensor* shift = ggml_reshape_4d(ctx->ggml_ctx, ggml_ext_slice(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); @@ -236,7 +236,7 @@ namespace H3VUpsampler { 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"); + w = ctx->weight_adapter->patch_weight(ctx, w, prefix + "weight"); } ggml_tensor* w32 = ggml_cast(ctx->ggml_ctx, w, GGML_TYPE_F32); ggml_tensor* xa = x; @@ -250,7 +250,7 @@ namespace H3VUpsampler { // 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* zero = ggml_scale(ctx->ggml_ctx, ggml_ext_slice(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); @@ -261,15 +261,15 @@ namespace H3VUpsampler { 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* xs = ggml_ext_slice(ctx, xp, 2, j, j + T); + ggml_tensor* wj = ggml_ext_slice(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"); + b = ctx->weight_adapter->patch_weight(ctx, b, prefix + "bias"); } acc = ggml_add(ctx->ggml_ctx, acc, ggml_reshape_4d(ctx->ggml_ctx, b, 1, 1, 1, channels)); } diff --git a/src/model/upscaler/ltx_latent_upscaler.hpp b/src/model/upscaler/ltx_latent_upscaler.hpp index 6650adc70b..f7910cce86 100644 --- a/src/model/upscaler/ltx_latent_upscaler.hpp +++ b/src/model/upscaler/ltx_latent_upscaler.hpp @@ -158,15 +158,15 @@ namespace LTXVUpsampler { const int64_t W = x->ne[0]; const int64_t H = x->ne[1]; const int64_t T = x->ne[2]; - x = ggml_ext_cont(ctx->ggml_ctx, x); + x = ggml_ext_cont(ctx, x); x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, num_channels, 1); x = ggml_group_norm(ctx->ggml_ctx, x, num_groups, eps); ggml_tensor* weight = params["weight"]; ggml_tensor* bias = params["bias"]; if (ctx->weight_adapter) { - weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight"); - bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias"); + weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight"); + bias = ctx->weight_adapter->patch_weight(ctx, bias, prefix + "bias"); } weight = ggml_reshape_4d(ctx->ggml_ctx, weight, 1, 1, num_channels, 1); bias = ggml_reshape_4d(ctx->ggml_ctx, bias, 1, 1, num_channels, 1); @@ -218,9 +218,9 @@ namespace LTXVUpsampler { int64_t w = x->ne[0]; GGML_ASSERT(x->ne[2] % (upscale_factor * upscale_factor) == 0); // x: [b*f, c*p1*p2, h, w] -> [b*f, c, h*p1, w*p2] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [b*f, h, w, c*p1*p2] - x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [b*f, h*w, c*p1*p2] - return DiT::unpatchify(ctx->ggml_ctx, x, h, w, upscale_factor, upscale_factor, true); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [b*f, h, w, c*p1*p2] + x = ggml_reshape_3d(ctx->ggml_ctx, x, x->ne[0], x->ne[1] * x->ne[2], x->ne[3]); // [b*f, h*w, c*p1*p2] + return DiT::unpatchify(ctx, x, h, w, upscale_factor, upscale_factor, true); } }; @@ -241,9 +241,9 @@ namespace LTXVUpsampler { const int64_t C = x->ne[3] / upscale_factor; // x: [b, c*p, f, h, w] -> [b, c, f*p, h, w] - x = ggml_ext_cont(ctx->ggml_ctx, x); + x = ggml_ext_cont(ctx, x); x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, F, upscale_factor, C); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, F * upscale_factor, C); } }; @@ -324,11 +324,11 @@ namespace LTXVUpsampler { auto blur_down = std::dynamic_pointer_cast(blocks["blur_down"]); // rearrange(x, "b c f h w -> (b f) c h w") - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); x = conv->forward(ctx, x); x = pixel_shuffle->forward(ctx, x); x = blur_down->forward(ctx, x); - return ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + return ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); } }; @@ -403,18 +403,18 @@ namespace LTXVUpsampler { } else if (config.temporal_upsample) { auto upsample_conv = std::dynamic_pointer_cast(blocks["upsampler.0"]); auto pixel_shuffle = std::dynamic_pointer_cast(blocks["upsampler.1"]); - x = upsample_conv->forward(ctx, x); // [b, c*2, f, h, w] - x = pixel_shuffle->forward(ctx, x); // [b, c, f*2, h, w] - x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); // x[:, :, 1:, :, :] + x = upsample_conv->forward(ctx, x); // [b, c*2, f, h, w] + x = pixel_shuffle->forward(ctx, x); // [b, c, f*2, h, w] + x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]); // x[:, :, 1:, :, :] } else { auto upsample_conv = std::dynamic_pointer_cast(blocks["upsampler.0"]); auto pixel_shuffle = std::dynamic_pointer_cast(blocks["upsampler.1"]); // rearrange(x, "b c f h w -> (b f) c h w"), - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*f, c, h, w] - x = upsample_conv->forward(ctx, x); // [b*f, c*4, h, w] - x = pixel_shuffle->forward(ctx, x); // [b*f, c, h*2, w*2] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [b*c, f, h, w] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // [b*f, c, h, w] + x = upsample_conv->forward(ctx, x); // [b*f, c*4, h, w] + x = pixel_shuffle->forward(ctx, x); // [b*f, c, h*2, w*2] + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // [b*c, f, h, w] } sd::ggml_graph_cut::mark_graph_cut(x, "ltx_latent_upsampler.spatial_up", "x"); diff --git a/src/model/upscaler/mosr.hpp b/src/model/upscaler/mosr.hpp index 0aa158c81c..6870e33d30 100644 --- a/src/model/upscaler/mosr.hpp +++ b/src/model/upscaler/mosr.hpp @@ -118,9 +118,9 @@ class MoSRLayerNorm2d : public UnaryBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { // [N, C, H, W] -> [N, H, W, C] so the norm reduces over channels. - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); - x = ggml_ext_layer_norm(ctx->ggml_ctx, x, params["weight"], params["bias"], 1e-6f); - return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); + x = ggml_ext_layer_norm(ctx, x, params["weight"], params["bias"], 1e-6f); + return ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); } }; @@ -141,9 +141,8 @@ class MoSRDepthwiseConv2d : public UnaryBlock { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { const int pad = kernel_size / 2; - x = ggml_conv_2d_dw_direct(ctx->ggml_ctx, params["weight"], ggml_ext_cont(ctx->ggml_ctx, x), 1, 1, pad, pad, 1, 1); - auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, channels, 1); - return ggml_add_inplace(ctx->ggml_ctx, x, b); + return ggml_ext_conv_2d_grouped(ctx, ggml_ext_cont(ctx, x), params["weight"], params["bias"], channels, + 1, 1, pad, pad, 1, 1, true); } }; @@ -179,10 +178,10 @@ class MoSRGatedCNNBlock : public GGMLBlock { auto i = channel_view(hidden, hidden - conv_ch); auto c = conv->forward(ctx, channel_view(2 * hidden - conv_ch, conv_ch)); if (hidden > conv_ch) { - c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, i), c, 2); + c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx, i), c, 2); } - h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, g))); + h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx, g))); h = mosr_mish(ctx->ggml_ctx, fc2->forward(ctx, h)); return ggml_add(ctx->ggml_ctx, h, ggml_scale_bias(ctx->ggml_ctx, x, 1.f, -0.5f)); } diff --git a/src/model/vae/auto_encoder_kl.hpp b/src/model/vae/auto_encoder_kl.hpp index 3a4733879a..3199377eb8 100644 --- a/src/model/vae/auto_encoder_kl.hpp +++ b/src/model/vae/auto_encoder_kl.hpp @@ -216,7 +216,7 @@ class VideoResnetBlock : public ResnetBlock { } ggml_tensor* get_alpha(GGMLRunnerContext* ctx) { - auto mix_factor = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["mix_factor"]); + auto mix_factor = ggml_ext_cast_f32(ctx, params["mix_factor"]); return ggml_sigmoid(ctx->ggml_ctx, mix_factor); } @@ -601,11 +601,11 @@ class AutoEncoderKLModel : public GGMLBlock { int64_t H = h * p; int64_t W = w * p; - z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N); // [N, C, p*p, h*w] - z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 1, 0, 2, 3)); // [N, C, h*w, p*p] - z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N); // [N*C*h, w, p, p] - z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, z, 0, 2, 1, 3)); // [N*C*h, p, w, p] - z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N); // [N, C, h*p, w*p] + z = ggml_reshape_4d(ctx->ggml_ctx, z, w * h, p * p, C, N); // [N, C, p*p, h*w] + z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, z, 1, 0, 2, 3)); // [N, C, h*w, p*p] + z = ggml_reshape_4d(ctx->ggml_ctx, z, p, p, w, h * C * N); // [N*C*h, w, p, p] + z = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, z, 0, 2, 1, 3)); // [N*C*h, p, w, p] + z = ggml_reshape_4d(ctx->ggml_ctx, z, W, H, C, N); // [N, C, h*p, w*p] } if (use_quant) { @@ -632,7 +632,7 @@ class AutoEncoderKLModel : public GGMLBlock { // sd::ggml_graph_cut::mark_graph_cut(z, "vae.encode.final", "z"); } if (sd_version_uses_flux2_vae(version)) { - z = ggml_ext_chunk(ctx->ggml_ctx, z, 2, 2)[0]; + z = ggml_ext_chunk(ctx, z, 2, 2)[0]; // [N, C, H, W] -> [N, C*p*p, H/p, W/p] int64_t p = 2; diff --git a/src/model/vae/hunyuan_vae.hpp b/src/model/vae/hunyuan_vae.hpp index 1aab8ac53c..ff2b2bb63f 100644 --- a/src/model/vae/hunyuan_vae.hpp +++ b/src/model/vae/hunyuan_vae.hpp @@ -96,20 +96,20 @@ namespace Hunyuan { auto previous = carry->take(); GGML_ASSERT(previous->ne[2] <= pad_t); for (int64_t frame = 0; frame < previous->ne[2]; frame++) { - temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, previous, 2, frame, frame + 1)); + temporal_frames.push_back(ggml_ext_slice(ctx, previous, 2, frame, frame + 1)); } for (int64_t frame = previous->ne[2]; frame < pad_t; frame++) { - temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1)); + temporal_frames.push_back(ggml_ext_slice(ctx, x, 2, 0, 1)); } } else { - auto first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto first = ggml_ext_slice(ctx, x, 2, 0, 1); for (int frame = 0; frame < pad_t; frame++) { temporal_frames.push_back(first); } } } for (int64_t frame = 0; frame < x->ne[2]; frame++) { - temporal_frames.push_back(ggml_ext_slice(ctx->ggml_ctx, x, 2, frame, frame + 1)); + temporal_frames.push_back(ggml_ext_slice(ctx, x, 2, frame, frame + 1)); } if (pad_t > 0 && carry != nullptr && carry->output != nullptr) { @@ -127,13 +127,13 @@ namespace Hunyuan { } auto replicate_pad = [&](ggml_tensor* input, int dim, int left, int right) { if (left > 0) { - auto first = ggml_ext_slice(ctx->ggml_ctx, input, dim, 0, 1); + auto first = ggml_ext_slice(ctx, input, dim, 0, 1); for (int i = 0; i < left; i++) { input = ggml_concat(ctx->ggml_ctx, first, input, dim); } } if (right > 0) { - auto last = ggml_ext_slice(ctx->ggml_ctx, input, dim, input->ne[dim] - 1, input->ne[dim]); + auto last = ggml_ext_slice(ctx, input, dim, input->ne[dim] - 1, input->ne[dim]); for (int i = 0; i < right; i++) { input = ggml_concat(ctx->ggml_ctx, input, last, dim); } @@ -184,19 +184,19 @@ namespace Hunyuan { auto k = k_proj->forward(ctx, x); // [b*c, t, h, w] auto v = v_proj->forward(ctx, x); // [b*c, t, h, w] - q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b); // [b, c, t*h*w] - q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 1, 0, 2, 3)); // [b, t*h*w, c] + q = ggml_reshape_3d(ctx->ggml_ctx, q, w * h * t, c, b); // [b, c, t*h*w] + q = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 1, 0, 2, 3)); // [b, t*h*w, c] - k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b); // [b, c, t*h*w] - k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 1, 0, 2, 3)); // [b, t*h*w, c] + k = ggml_reshape_3d(ctx->ggml_ctx, k, w * h * t, c, b); // [b, c, t*h*w] + k = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, k, 1, 0, 2, 3)); // [b, t*h*w, c] - v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w] - v = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c] + v = ggml_reshape_3d(ctx->ggml_ctx, v, w * h * t, c, b); // [b, c, t*h*w] + v = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, v, 1, 0, 2, 3)); // [b, t*h*w, c] x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [b, t*h*w, c] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w] - x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w] + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [b, c, t*h*w] + x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, t, c * b); // [b*c, t, h, w] x = proj_out->forward(ctx, x); @@ -301,13 +301,13 @@ namespace Hunyuan { int64_t W = x->ne[0]; x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, C, factor); // [factor, C, T*H, W] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // [C, factor, T*H, W] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // [C, factor, T*H, W] x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C); // [C*factor_t*factor_s, factor_s, T*H, W] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C); // [C*factor_t, factor_s, T*H, W*factor_s] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C); // [C, factor_t, T, H*factor_s*W*factor_s] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C); // [C, T*factor_t, H*factor_s, W*factor_s] return x; @@ -328,11 +328,11 @@ namespace Hunyuan { ggml_tensor* shortcut = nullptr; if (factor_t == 2 && !continuation) { - auto h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1); // [B*factor*OC, 1, H, W] - h_first = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B); // [B*2*OC, 1, H*factor_s, W*factor_s] - h_first = ggml_ext_slice(ctx->ggml_ctx, h_first, 3, 0, out_channels); // [B*OC, 1, H*factor_s, W*factor_s] + auto h_first = ggml_ext_slice(ctx, h, 2, 0, 1); // [B*factor*OC, 1, H, W] + h_first = _pixel_shuffle_3d(ctx, h_first, 1, factor_s, B); // [B*2*OC, 1, H*factor_s, W*factor_s] + h_first = ggml_ext_slice(ctx, h_first, 3, 0, out_channels); // [B*OC, 1, H*factor_s, W*factor_s] - auto x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto x_first = ggml_ext_slice(ctx, x, 2, 0, 1); x_first = repeat_interleave_channels(ctx, x_first, repeats / 2, x->ne[0], x->ne[1], 1); x_first = _pixel_shuffle_3d(ctx, x_first, 1, factor_s, B); @@ -340,12 +340,12 @@ namespace Hunyuan { return ggml_add(ctx->ggml_ctx, h_first, x_first); } - auto h_next = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]); // [B*factor*OC, T - 1, H, W] + auto h_next = ggml_ext_slice(ctx, h, 2, 1, h->ne[2]); // [B*factor*OC, T - 1, H, W] h_next = _pixel_shuffle_3d(ctx, h_next, factor_t, factor_s, B); // [B*OC, (T - 1)*factor_t, H*factor_s, W*factor_s] h = ggml_concat(ctx->ggml_ctx, h_first, h_next, 2); // [B*OC, 1 + (T - 1)*factor_t, H*factor_s, W*factor_s] - auto x_next = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); + auto x_next = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]); x_next = repeat_interleave_channels(ctx, x_next, repeats, x->ne[0], x->ne[1], x->ne[2] - 1); x_next = _pixel_shuffle_3d(ctx, x_next, factor_t, factor_s, B); @@ -374,13 +374,13 @@ namespace Hunyuan { int64_t factor = factor_t * factor_s * factor_s; x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, factor_t, T, C); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s, H * T, factor_t * C); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W, H * T, factor_s * factor_t * C); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor, C); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); return ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, C * factor); } @@ -398,9 +398,9 @@ namespace Hunyuan { int64_t spatial = W * H * T; int64_t groups = x->ne[3] / group_size; x = ggml_reshape_3d(ctx->ggml_ctx, x, spatial, group_size, groups); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); x = ggml_sum_rows(ctx->ggml_ctx, x); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H, T, groups); return ggml_scale(ctx->ggml_ctx, x, 1.f / static_cast(group_size)); } @@ -435,19 +435,19 @@ namespace Hunyuan { ggml_tensor* h_first = nullptr; ggml_tensor* x_first = nullptr; if (factor_t == 2) { - h_first = ggml_ext_slice(ctx->ggml_ctx, h, 2, 0, 1); + h_first = ggml_ext_slice(ctx, h, 2, 0, 1); h_first = pixel_unshuffle_3d(ctx, h_first, 1, factor_s); h_first = ggml_concat(ctx->ggml_ctx, h_first, h_first, 3); - x_first = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + x_first = ggml_ext_slice(ctx, x, 2, 0, 1); x_first = pixel_unshuffle_3d(ctx, x_first, 1, factor_s); x_first = mean_channel_groups(ctx, x_first, group_size / 2); if (x->ne[2] == 1) { return ggml_add(ctx->ggml_ctx, h_first, x_first); } - h = ggml_ext_slice(ctx->ggml_ctx, h, 2, 1, h->ne[2]); - x = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1, x->ne[2]); + h = ggml_ext_slice(ctx, h, 2, 1, h->ne[2]); + x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]); } GGML_ASSERT(h->ne[2] % factor_t == 0); @@ -648,7 +648,7 @@ namespace Hunyuan { x = ggml_silu_inplace(ctx->ggml_ctx, x); x = conv_out->forward(ctx, x); x = ggml_add(ctx->ggml_ctx, x, shortcut); - return ggml_ext_slice(ctx->ggml_ctx, x, 3, 0, z_channels); + return ggml_ext_slice(ctx, x, 3, 0, z_channels); } }; @@ -707,7 +707,7 @@ namespace Hunyuan { const int64_t frames = h->ne[2]; for (int64_t start = 0; start < frames; start += HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE) { const int64_t end = std::min(start + HUNYUAN_VIDEO_VAE_TEMPORAL_CHUNK_SIZE, frames); - auto chunk = ggml_ext_slice(ctx->ggml_ctx, h, 2, start, end); + auto chunk = ggml_ext_slice(ctx, h, 2, start, end); std::vector carry_output; TemporalConvCarry carry{ diff --git a/src/model/vae/ltx_audio_vae.hpp b/src/model/vae/ltx_audio_vae.hpp index a05ce86a46..631b8c5a4f 100644 --- a/src/model/vae/ltx_audio_vae.hpp +++ b/src/model/vae/ltx_audio_vae.hpp @@ -220,7 +220,7 @@ namespace LTXV { auto x = ggml_reshape_3d(ctx, waveform, time, 1, channels * batch); if (left_pad > 0) { - x = ggml_ext_pad_ext(ctx, runner_ctx->backend, x, static_cast(left_pad), 0, 0, 0, 0, 0, 0, 0); + x = ggml_ext_pad_ext(runner_ctx, x, static_cast(left_pad), 0, 0, 0, 0, 0, 0, 0); } auto frames = ggml_conv_1d(ctx, forward_basis, x, hop_length, 0, 1); @@ -228,8 +228,8 @@ namespace LTXV { GGML_ASSERT(frames->ne[1] == stft_channels); GGML_ASSERT(frames->ne[2] == channels * batch); - auto real = ggml_ext_slice(ctx, frames, 1, 0, n_freqs); - auto imag = ggml_ext_slice(ctx, frames, 1, n_freqs, stft_channels); + auto real = ggml_ext_slice(runner_ctx, frames, 1, 0, n_freqs); + auto imag = ggml_ext_slice(runner_ctx, frames, 1, n_freqs, stft_channels); auto magnitude = ggml_sqrt(ctx, ggml_add(ctx, ggml_sqr(ctx, real), @@ -265,42 +265,19 @@ namespace LTXV { return type == GGML_TYPE_BF16 ? GGML_TYPE_F16 : type; } - static ggml_tensor* repeat_with_vulkan_f32_workaround(ggml_backend_t backend, - ggml_context* ctx, - ggml_tensor* x, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - if (x->type != GGML_TYPE_F32 && - (x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && - sd_backend_is(backend, "vulkan")) { - auto x_f32 = ggml_cast(ctx, x, GGML_TYPE_F32); - auto repeated = ggml_repeat_4d(ctx, - x_f32, - ne0, - ne1, - ne2, - ne3); - return ggml_cast(ctx, repeated, x->type); - } - return ggml_repeat_4d(ctx, x, ne0, ne1, ne2, ne3); - } - static ggml_tensor* repeat_1d_value(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t count) { - auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(x->ne[0] == 1); - return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, x, count, x->ne[1], x->ne[2], x->ne[3]); + return ggml_ext_repeat_4d(runner_ctx, x, count, x->ne[1], x->ne[2], x->ne[3]); } static ggml_tensor* replicate_pad_1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t left, int64_t right) { auto ctx = runner_ctx->ggml_ctx; if (left > 0) { - auto first = ggml_ext_slice(ctx, x, 0, 0, 1); + auto first = ggml_ext_slice(runner_ctx, x, 0, 0, 1); x = ggml_concat(ctx, repeat_1d_value(runner_ctx, first, left), x, 0); } if (right > 0) { - auto last = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]); + auto last = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]); x = ggml_concat(ctx, x, repeat_1d_value(runner_ctx, last, right), 0); } return x; @@ -314,7 +291,7 @@ namespace LTXV { } else if (ggml_n_dims(base) == 1) { base = ggml_reshape_4d(ctx, base, base->ne[0], 1, 1, 1); } - return repeat_with_vulkan_f32_workaround(runner_ctx->backend, ctx, base, base->ne[0], 1, channels, 1); + return ggml_ext_repeat_4d(runner_ctx, base, base->ne[0], 1, channels, 1); } static ggml_tensor* depthwise_conv1d(GGMLRunnerContext* runner_ctx, @@ -329,7 +306,8 @@ namespace LTXV { return ggml_reshape_4d(ctx, out, out->ne[0], out->ne[1], 1, 1); } - static ggml_tensor* reverse_1d_filter(ggml_context* ctx, ggml_tensor* filter) { + static ggml_tensor* reverse_1d_filter(GGMLRunnerContext* runner_ctx, ggml_tensor* filter) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(ctx != nullptr); GGML_ASSERT(filter != nullptr); GGML_ASSERT(filter->ne[1] == 1); @@ -338,16 +316,17 @@ namespace LTXV { ggml_tensor* reversed = nullptr; for (int64_t k = filter->ne[0] - 1; k >= 0; --k) { - auto slice = ggml_ext_slice(ctx, filter, 0, k, k + 1); + auto slice = ggml_ext_slice(runner_ctx, filter, 0, k, k + 1); reversed = reversed == nullptr ? slice : ggml_concat(ctx, reversed, slice, 0); } return reversed; } - static ggml_tensor* depthwise_conv_transpose1d(ggml_context* ctx, + static ggml_tensor* depthwise_conv_transpose1d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* filter, int stride) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(x->ne[2] == 1 && x->ne[3] == 1); GGML_ASSERT(filter->ne[1] == 1); GGML_ASSERT(filter->ne[2] == 1 && filter->ne[3] == 1); @@ -359,7 +338,7 @@ namespace LTXV { auto x_flat = ggml_reshape_3d(ctx, x, 1, time, channels); if (stride > 1) { - auto zero_unit = ggml_ext_scale(ctx, x_flat, 0.0f); + auto zero_unit = ggml_ext_scale(runner_ctx, x_flat, 0.0f); auto zero_tail = zero_unit; for (int i = 1; i < stride - 1; ++i) { zero_tail = ggml_concat(ctx, zero_tail, zero_unit, 0); @@ -368,16 +347,16 @@ namespace LTXV { } x_flat = ggml_reshape_3d(ctx, x_flat, time * stride, 1, channels); - auto reversed_filter = reverse_1d_filter(ctx, filter); + auto reversed_filter = reverse_1d_filter(runner_ctx, filter); auto out = ggml_conv_1d(ctx, reversed_filter, x_flat, 1, static_cast(kernel_size - 1), 1); if (out->ne[0] > out_time) { - out = ggml_ext_slice(ctx, out, 0, 0, out_time); + out = ggml_ext_slice(runner_ctx, out, 0, 0, out_time); } GGML_ASSERT(out->ne[0] == out_time); GGML_ASSERT(out->ne[1] == 1); GGML_ASSERT(out->ne[2] == channels); - out = ggml_ext_scale(ctx, out, static_cast(stride)); + out = ggml_ext_scale(runner_ctx, out, static_cast(stride)); return ggml_reshape_4d(ctx, out, out_time, channels, 1, 1); } @@ -409,21 +388,22 @@ namespace LTXV { auto x = ggml_reshape_3d(ctx, waveform, time, channels * batch, 1); x = replicate_pad_1d(runner_ctx, x, pad, pad); - x = depthwise_conv_transpose1d(ctx, x, filter, ratio); - x = ggml_ext_slice(ctx, x, 0, pad_left, x->ne[0] - pad_right); + x = depthwise_conv_transpose1d(runner_ctx, x, filter, ratio); + x = ggml_ext_slice(runner_ctx, x, 0, pad_left, x->ne[0] - pad_right); return ggml_reshape_3d(ctx, x, x->ne[0], channels, batch); } - static ggml_tensor* crop_waveform_samples(ggml_context* ctx, + static ggml_tensor* crop_waveform_samples(GGMLRunnerContext* runner_ctx, ggml_tensor* waveform, int64_t target_samples) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(ctx != nullptr); GGML_ASSERT(waveform != nullptr); if (waveform->ne[0] == target_samples) { return waveform; } GGML_ASSERT(waveform->ne[0] > target_samples); - return ggml_ext_slice(ctx, waveform, 0, 0, target_samples); + return ggml_ext_slice(runner_ctx, waveform, 0, 0, target_samples); } struct PixelNorm2D : public UnaryBlock { @@ -433,9 +413,9 @@ namespace LTXV { : eps(eps) {} ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); h = ggml_rms_norm(ctx->ggml_ctx, h, eps); - h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 0, 3)); + h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 0, 3)); return h; } }; @@ -456,8 +436,7 @@ namespace LTXV { auto conv = std::dynamic_pointer_cast(blocks["conv"]); int pad_h = kernel_size.first - 1; int pad_w = kernel_size.second - 1; - x = ggml_ext_pad_ext(ctx->ggml_ctx, - ctx->backend, + x = ggml_ext_pad_ext(ctx, x, pad_w / 2, pad_w - pad_w / 2, @@ -481,7 +460,7 @@ namespace LTXV { auto conv = std::dynamic_pointer_cast(blocks["conv"]); x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST); x = conv->forward(ctx, x); - return ggml_ext_slice(ctx->ggml_ctx, x, 1, 1, x->ne[1]); + return ggml_ext_slice(ctx, x, 1, 1, x->ne[1]); } }; @@ -606,7 +585,7 @@ namespace LTXV { GGML_ASSERT(dilation == 1); x = ggml_conv_transpose_1d(ctx->ggml_ctx, params["weight"], x, stride, 0, dilation); if (padding > 0) { - x = ggml_ext_slice(ctx->ggml_ctx, x, 0, padding, x->ne[0] - padding); + x = ggml_ext_slice(ctx, x, 0, padding, x->ne[0] - padding); } if (bias) { auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, params["bias"]->ne[0], 1, 1); @@ -639,7 +618,7 @@ namespace LTXV { beta = ggml_reshape_4d(ctx->ggml_ctx, beta, 1, beta->ne[0], 1, 1); auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha)); oscillation = ggml_mul(ctx->ggml_ctx, oscillation, oscillation); - auto eps_tensor = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), eps); + auto eps_tensor = ggml_ext_scale(ctx, ggml_ext_ones(ctx, 1, 1, 1, 1), eps); oscillation = ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, beta, eps_tensor)); return ggml_add(ctx->ggml_ctx, x, oscillation); } @@ -675,8 +654,8 @@ namespace LTXV { int up_pad_right = up_pad * up_ratio + (up_kernel_size - up_ratio + 1) / 2; x = replicate_pad_1d(ctx, x, up_pad, up_pad); - x = depthwise_conv_transpose1d(ctx->ggml_ctx, x, up_filter, up_ratio); - x = ggml_ext_slice(ctx->ggml_ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); + x = depthwise_conv_transpose1d(ctx, x, up_filter, up_ratio); + x = ggml_ext_slice(ctx, x, 0, up_pad_left, x->ne[0] - up_pad_right); x = act->forward(ctx, x); @@ -778,7 +757,7 @@ namespace LTXV { ggml_tensor* prepare_input(GGMLRunnerContext* ctx, ggml_tensor* mel) { mel = ggml_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, mel, 1, 0, 2, 3)); - auto mels = ggml_ext_chunk(ctx->ggml_ctx, mel, 2, 2); + auto mels = ggml_ext_chunk(ctx, mel, 2, 2); mel = ggml_concat(ctx->ggml_ctx, mels[0], mels[1], 1); // mel = ggml_reshape_4d(ctx->ggml_ctx, // mel, @@ -810,7 +789,7 @@ namespace LTXV { auto block_out = resblock->forward(ctx, x); sum = sum == nullptr ? block_out : ggml_add(ctx->ggml_ctx, sum, block_out); } - x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.0f / static_cast(resblock_kernel_sizes.size())); + x = ggml_ext_scale(ctx, sum, 1.0f / static_cast(resblock_kernel_sizes.size())); } x = act_post->forward(ctx, x); @@ -876,8 +855,8 @@ namespace LTXV { int target_freq) { int64_t time = std::min(decoded->ne[1], target_time); int64_t freq = std::min(decoded->ne[0], target_freq); - decoded = ggml_ext_slice(ctx->ggml_ctx, decoded, 0, 0, freq); - decoded = ggml_ext_slice(ctx->ggml_ctx, decoded, 1, 0, time); + decoded = ggml_ext_slice(ctx, decoded, 0, 0, freq); + decoded = ggml_ext_slice(ctx, decoded, 1, 0, time); return decoded; } @@ -995,7 +974,7 @@ namespace LTXV { ggml_add(ctx->ggml_ctx, residual, skip), -1.0f, 1.0f); - waveform = crop_waveform_samples(ctx->ggml_ctx, waveform, out_time); + waveform = crop_waveform_samples(ctx, waveform, out_time); } return waveform; diff --git a/src/model/vae/ltx_vae.hpp b/src/model/vae/ltx_vae.hpp index 3fa42a1662..743549be5b 100644 --- a/src/model/vae/ltx_vae.hpp +++ b/src/model/vae/ltx_vae.hpp @@ -32,28 +32,26 @@ namespace LTXVAE { return ggml_reshape_4d(ctx, x, 1, 1, 1, ggml_nelements(x)); } - static inline std::pair get_shift_scale(ggml_context* ctx, - ggml_tensor* table, - ggml_tensor* timestep, - int64_t channels, - int parts) { + static inline std::pair get_shift_scale(GGMLRunnerContext* runner_ctx, ggml_tensor* table, ggml_tensor* timestep, int64_t channels, int parts) { + auto ctx = runner_ctx->ggml_ctx; GGML_ASSERT(timestep != nullptr); GGML_ASSERT(ggml_nelements(timestep) == channels * parts); auto timestep_view = ggml_reshape_2d(ctx, timestep, channels, parts); auto values = ggml_add(ctx, table, timestep_view); - auto chunks = ggml_ext_chunk(ctx, values, parts, 1, false); + auto chunks = ggml_ext_chunk(runner_ctx, values, parts, 1, false); auto shift = reshape_channel_broadcast(ctx, ggml_cont(ctx, chunks[0])); auto scale = reshape_channel_broadcast(ctx, ggml_cont(ctx, chunks[1])); return {shift, scale}; } - static inline ggml_tensor* depth_to_space_3d(ggml_context* ctx, + static inline ggml_tensor* depth_to_space_3d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t c, int factor_t, int factor_s, bool drop_first_temporal_frame) { + auto ctx = runner_ctx->ggml_ctx; // x: [B*c*p1*p2*p3, T, H, W], B == 1, p2 == p3 == factor_s, p1 == factor_t // return: [B*c, T*p1, H*p2, W*p2] // Match: rearrange(x, "b (c p1 p2 p3) d h w -> b c (d p1) (h p2) (w p3)") @@ -61,26 +59,24 @@ namespace LTXVAE { const int64_t H = x->ne[1]; const int64_t W = x->ne[0]; - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // [T, C, H, W] - x = ggml_reshape_4d(ctx, x, W, H, factor_s, factor_s * factor_t * c * T); // [T*c*p1*p2, p3, H, W] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [T*c*p1*p2, H, W, p3] - x = ggml_reshape_4d(ctx, x, factor_s * W, H, factor_s, factor_t * c * T); // [T*c*p1, p2, H, W*p3] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [T*c*p1, H, p2, W*p3] - x = ggml_reshape_4d(ctx, x, factor_s * W * factor_s * H, factor_t, c, T); // [T, c, p1, H*p2*W*p3] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // [c, T, p1, H*p2*W*p3] - x = ggml_reshape_4d(ctx, x, factor_s * W, factor_s * H, factor_t * T, c); // [T, c, T*p1, H*p2*W*p3] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 1, 3, 2)); // [T, C, H, W] + x = ggml_reshape_4d(ctx, x, W, H, factor_s, factor_s * factor_t * c * T); // [T*c*p1*p2, p3, H, W] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3)); // [T*c*p1*p2, H, W, p3] + x = ggml_reshape_4d(ctx, x, factor_s * W, H, factor_s, factor_t * c * T); // [T*c*p1, p2, H, W*p3] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [T*c*p1, H, p2, W*p3] + x = ggml_reshape_4d(ctx, x, factor_s * W * factor_s * H, factor_t, c, T); // [T, c, p1, H*p2*W*p3] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 1, 3, 2)); // [c, T, p1, H*p2*W*p3] + x = ggml_reshape_4d(ctx, x, factor_s * W, factor_s * H, factor_t * T, c); // [T, c, T*p1, H*p2*W*p3] if (drop_first_temporal_frame && factor_t > 1 && x->ne[2] > 0) { - x = ggml_ext_slice(ctx, x, 2, 1, x->ne[2]); + x = ggml_ext_slice(runner_ctx, x, 2, 1, x->ne[2]); } return x; } - static inline ggml_tensor* patchify(ggml_context* ctx, - ggml_tensor* x, - int patch_size) { - return WAN::WanVAE::patchify(ctx, x, patch_size, 1); + static inline ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int patch_size) { + return WAN::WanVAE::patchify(runner_ctx, x, patch_size, 1); } class CausalConv3d : public GGMLBlock { @@ -113,20 +109,20 @@ namespace LTXVAE { auto conv = std::dynamic_pointer_cast(blocks["conv"]); if (causal) { - auto first_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto first_frame = ggml_ext_slice(ctx, x, 2, 0, 1); auto first_frame_pad = first_frame; for (int i = 1; i < time_kernel_size - 1; i++) { first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2); } x = ggml_concat(ctx->ggml_ctx, first_frame_pad, x, 2); } else { - auto first_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto first_frame = ggml_ext_slice(ctx, x, 2, 0, 1); auto first_frame_pad = first_frame; for (int i = 1; i < (time_kernel_size - 1) / 2; i++) { first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2); } - auto last_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - 1, x->ne[2]); + auto last_frame = ggml_ext_slice(ctx, x, 2, x->ne[2] - 1, x->ne[2]); auto last_frame_pad = last_frame; for (int i = 1; i < (time_kernel_size - 1) / 2; i++) { last_frame_pad = ggml_concat(ctx->ggml_ctx, last_frame_pad, last_frame, 2); @@ -165,7 +161,7 @@ namespace LTXVAE { GGML_ASSERT(start_idx >= 0); GGML_ASSERT(end_idx > 0); - auto slice = ggml_ext_slice(ctx->ggml_ctx, x, 2, start_idx, end_idx); + auto slice = ggml_ext_slice(ctx, x, 2, start_idx, end_idx); feat_map[feat_idx] = ggml_cont(ctx->ggml_ctx, slice); } feat_idx++; @@ -175,7 +171,7 @@ namespace LTXVAE { if (prev != nullptr) { left_pad = prev; } else { - auto first_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto first_frame = ggml_ext_slice(ctx, x, 2, 0, 1); left_pad = first_frame; for (int i = 1; i < pad; i++) { left_pad = ggml_concat(ctx->ggml_ctx, left_pad, first_frame, 2); @@ -185,7 +181,7 @@ namespace LTXVAE { } if (!causal && pad > 0) { - auto last_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, x->ne[2] - 1, x->ne[2]); + auto last_frame = ggml_ext_slice(ctx, x, 2, x->ne[2] - 1, x->ne[2]); auto right_pad = last_frame; for (int i = 1; i < pad; i++) { right_pad = ggml_concat(ctx->ggml_ctx, right_pad, last_frame, 2); @@ -204,9 +200,9 @@ namespace LTXVAE { : eps(eps) {} ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { - auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2)); + auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 3, 0, 1, 2)); h = ggml_rms_norm(ctx->ggml_ctx, h, eps); - h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0)); + h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 3, 0)); return h; } }; @@ -267,7 +263,7 @@ namespace LTXVAE { auto values = ggml_add(ctx->ggml_ctx, params["scale_shift_table"], ggml_reshape_2d(ctx->ggml_ctx, timestep, channels, 4)); - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, values, 4, 1, false); + auto chunks = ggml_ext_chunk(ctx, values, 4, 1, false); shift1 = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[0])); scale1 = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[1])); shift2 = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[2])); @@ -313,7 +309,7 @@ namespace LTXVAE { auto values = ggml_add(ctx->ggml_ctx, params["scale_shift_table"], ggml_reshape_2d(ctx->ggml_ctx, timestep, channels, 4)); - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, values, 4, 1, false); + auto chunks = ggml_ext_chunk(ctx, values, 4, 1, false); shift1 = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[0])); scale1 = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[1])); shift2 = reshape_channel_broadcast(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, chunks[2])); @@ -430,7 +426,7 @@ namespace LTXVAE { ggml_tensor* x_in = nullptr; if (residual) { - x_in = depth_to_space_3d(ctx->ggml_ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, factor_t > 1); + x_in = depth_to_space_3d(ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, factor_t > 1); int repeat = (factor_t * factor_s * factor_s) / out_channels_reduction_factor; auto res = x_in; for (int i = 1; i < repeat; i++) { @@ -440,7 +436,7 @@ namespace LTXVAE { } x = conv->forward(ctx, x, causal); - x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, factor_t > 1); + x = depth_to_space_3d(ctx, x, get_output_channels(), factor_t, factor_s, factor_t > 1); if (residual) { x = ggml_add(ctx->ggml_ctx, x, x_in); } @@ -460,7 +456,7 @@ namespace LTXVAE { ggml_tensor* x_in = nullptr; if (residual) { - x_in = depth_to_space_3d(ctx->ggml_ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, drop_first); + x_in = depth_to_space_3d(ctx, x, in_channels / (factor_t * factor_s * factor_s), factor_t, factor_s, drop_first); int repeat = (factor_t * factor_s * factor_s) / out_channels_reduction_factor; auto res = x_in; for (int i = 1; i < repeat; i++) { @@ -470,7 +466,7 @@ namespace LTXVAE { } x = conv->forward(ctx, x, feat_map, feat_idx, chunk_idx, causal, temporal_pad); - x = depth_to_space_3d(ctx->ggml_ctx, x, get_output_channels(), factor_t, factor_s, drop_first); + x = depth_to_space_3d(ctx, x, get_output_channels(), factor_t, factor_s, drop_first); if (residual) { x = ggml_add(ctx->ggml_ctx, x, x_in); } @@ -515,7 +511,7 @@ namespace LTXVAE { auto conv_downsample = std::dynamic_pointer_cast(blocks["conv_downsample"]); if (factor_t > 1 && x->ne[2] > 0) { - auto first_frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); + auto first_frame = ggml_ext_slice(ctx, x, 2, 0, 1); auto first_frame_pad = first_frame; for (int i = 1; i < factor_t - 1; ++i) { first_frame_pad = ggml_concat(ctx->ggml_ctx, first_frame_pad, first_frame, 2); @@ -860,7 +856,7 @@ namespace LTXVAE { x = ggml_silu_inplace(ctx->ggml_ctx, x); x = conv_out->forward(ctx, x, true); - auto last_channel = ggml_ext_slice(ctx->ggml_ctx, x, 3, x->ne[3] - 1, x->ne[3]); + auto last_channel = ggml_ext_slice(ctx, x, 3, x->ne[3] - 1, x->ne[3]); auto repeat_shape = ggml_new_tensor_4d(ctx->ggml_ctx, last_channel->type, last_channel->ne[0], last_channel->ne[1], last_channel->ne[2], latent_channels - 1); auto repeated = ggml_repeat(ctx->ggml_ctx, last_channel, repeat_shape); return ggml_concat(ctx->ggml_ctx, x, repeated, 3); @@ -959,7 +955,7 @@ namespace LTXVAE { ggml_tensor* scaled_timestep = timestep; if (timestep_conditioning) { - auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]); + auto multiplier = ggml_ext_cast_f32(ctx, params["timestep_scale_multiplier"]); scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier); } @@ -981,7 +977,7 @@ namespace LTXVAE { if (timestep_conditioning) { auto last_time_embedder = std::dynamic_pointer_cast(blocks["last_time_embedder"]); auto timestep_embed = last_time_embedder->forward(ctx, scaled_timestep); - auto [shift, scale] = get_shift_scale(ctx->ggml_ctx, + auto [shift, scale] = get_shift_scale(ctx, params["last_scale_shift_table"], timestep_embed, hidden_channels, @@ -1010,7 +1006,7 @@ namespace LTXVAE { ggml_tensor* scaled_timestep = timestep; if (timestep_conditioning && timestep != nullptr) { - auto multiplier = ggml_ext_cast_f32(ctx->ggml_ctx, ctx->backend, params["timestep_scale_multiplier"]); + auto multiplier = ggml_ext_cast_f32(ctx, params["timestep_scale_multiplier"]); scaled_timestep = ggml_mul(ctx->ggml_ctx, timestep, multiplier); } @@ -1038,7 +1034,7 @@ namespace LTXVAE { if (timestep_conditioning) { auto last_time_embedder = std::dynamic_pointer_cast(blocks["last_time_embedder"]); auto timestep_embed = last_time_embedder->forward(ctx, scaled_timestep); - auto [shift, scale] = get_shift_scale(ctx->ggml_ctx, + auto [shift, scale] = get_shift_scale(ctx, params["last_scale_shift_table"], timestep_embed, hidden_channels, @@ -1092,7 +1088,7 @@ namespace LTXVAE { auto processor = std::dynamic_pointer_cast(blocks["per_channel_statistics"]); auto latents = processor->un_normalize(ctx, z); auto out = decoder->forward(ctx, latents, timestep); - out = WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1); + out = WAN::WanVAE::unpatchify(ctx, out, patch_size, 1); return out; } @@ -1112,7 +1108,7 @@ namespace LTXVAE { const int64_t T = z->ne[2]; if (T <= 1) { auto out = decoder->forward(ctx, latents, timestep); - return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1); + return WAN::WanVAE::unpatchify(ctx, out, patch_size, 1); } // feat_map holds ggml_tensor* nodes (contiguous copies at each conv layer). @@ -1145,20 +1141,20 @@ namespace LTXVAE { int chunk_overlap = overlap; // modified by forward_tiled_frame temporal inflation - auto z_chunk = ggml_ext_slice(ctx->ggml_ctx, latents, 2, i, end_i); + auto z_chunk = ggml_ext_slice(ctx, latents, 2, i, end_i); auto out_chunk = decoder->forward_tiled_frame(ctx, z_chunk, timestep, feat_map, feat_idx, i, chunk_overlap); // discard overlap frames if it's not the final chunk if (overlap > 0 && end_i < (int)T) { - out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap); + out_chunk = ggml_ext_slice(ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap); } out = (out == nullptr) ? out_chunk : ggml_concat(ctx->ggml_ctx, out, out_chunk, 2); } - return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out, patch_size, 1); + return WAN::WanVAE::unpatchify(ctx, out, patch_size, 1); } ggml_tensor* decode_tiled_chunk(GGMLRunnerContext* ctx, @@ -1177,9 +1173,9 @@ namespace LTXVAE { auto out_chunk = decoder->forward_tiled_frame(ctx, latents, timestep, feat_map, feat_idx, chunk_idx, chunk_overlap); if (chunk_overlap > 0) { - out_chunk = ggml_ext_slice(ctx->ggml_ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap); + out_chunk = ggml_ext_slice(ctx, out_chunk, 2, 0, out_chunk->ne[2] - chunk_overlap); } - return WAN::WanVAE::unpatchify(ctx->ggml_ctx, out_chunk, patch_size, 1); + return WAN::WanVAE::unpatchify(ctx, out_chunk, patch_size, 1); } ggml_tensor* encode(GGMLRunnerContext* ctx, @@ -1188,9 +1184,9 @@ namespace LTXVAE { auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); auto processor = std::dynamic_pointer_cast(blocks["per_channel_statistics"]); - x = patchify(ctx->ggml_ctx, x, patch_size); + x = patchify(ctx, x, patch_size); auto out = encoder->forward(ctx, x); - auto mean = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3, false)[0]; + auto mean = ggml_ext_chunk(ctx, out, 2, 3, false)[0]; mean = ggml_cont(ctx->ggml_ctx, mean); return processor->normalize(ctx, mean); } diff --git a/src/model/vae/mage_vae.hpp b/src/model/vae/mage_vae.hpp index 38ff89c4f0..838dfeaec0 100644 --- a/src/model/vae/mage_vae.hpp +++ b/src/model/vae/mage_vae.hpp @@ -32,23 +32,24 @@ namespace MageVAE { ggml_tensor* weight = affine ? params["weight"] : nullptr; ggml_tensor* bias = affine ? params["bias"] : nullptr; if (affine && ctx->weight_adapter) { - weight = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, weight, prefix + "weight"); - bias = ctx->weight_adapter->patch_weight(ctx->ggml_ctx, ctx->backend, bias, prefix + "bias"); + weight = ctx->weight_adapter->patch_weight(ctx, weight, prefix + "weight"); + bias = ctx->weight_adapter->patch_weight(ctx, bias, prefix + "bias"); } // [N, C, H, W] -> [N, H, W, C] so layer norm reduces over channels. - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); - x = ggml_ext_layer_norm(ctx->ggml_ctx, x, weight, bias, 1e-6f); - return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); + x = ggml_ext_layer_norm(ctx, x, weight, bias, 1e-6f); + return ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); } }; - inline ggml_tensor* modulate_2d(ggml_context* ctx, + inline ggml_tensor* modulate_2d(GGMLRunnerContext* runner_ctx, ggml_tensor* x, ggml_tensor* shift, ggml_tensor* scale) { - shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]); - scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]); - return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(ctx, 1, 1, 1, 1))), shift); + auto ctx = runner_ctx->ggml_ctx; + shift = ggml_reshape_4d(ctx, shift, 1, 1, shift->ne[0], shift->ne[1]); + scale = ggml_reshape_4d(ctx, scale, 1, 1, scale->ne[0], scale->ne[1]); + return ggml_add(ctx, ggml_mul(ctx, x, ggml_add(ctx, scale, ggml_ext_ones(runner_ctx, 1, 1, 1, 1))), shift); } inline ggml_tensor* channel_attention(GGMLRunnerContext* ctx, @@ -70,7 +71,7 @@ namespace MageVAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* timestep) { auto linear_0 = std::dynamic_pointer_cast(blocks["mlp.0"]); auto linear_2 = std::dynamic_pointer_cast(blocks["mlp.2"]); - auto x = ggml_ext_timestep_embedding(ctx->ggml_ctx, timestep, 256, 10000, 1.f); + auto x = ggml_ext_timestep_embedding(ctx, timestep, 256, 10000, 1.f); x = linear_0->forward(ctx, x); x = ggml_silu_inplace(ctx->ggml_ctx, x); return linear_2->forward(ctx, x); @@ -139,9 +140,9 @@ namespace MageVAE { auto ada = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"]); auto params = ada->forward(ctx, ggml_silu(ctx->ggml_ctx, condition)); - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 6, 0); + auto chunks = ggml_ext_chunk(ctx, params, 6, 0); auto x = norm1->forward(ctx, input); - x = modulate_2d(ctx->ggml_ctx, x, chunks[0], chunks[1]); + x = modulate_2d(ctx, x, chunks[0], chunks[1]); x = conv1->forward(ctx, x); x = conv2->forward(ctx, x); x = ggml_gelu(ctx->ggml_ctx, x); @@ -151,7 +152,7 @@ namespace MageVAE { x = ggml_add(ctx->ggml_ctx, input, ggml_mul(ctx->ggml_ctx, x, gate_1)); auto h = norm2->forward(ctx, x); - h = modulate_2d(ctx->ggml_ctx, h, chunks[3], chunks[4]); + h = modulate_2d(ctx, h, chunks[3], chunks[4]); h = conv4->forward(ctx, h); h = ggml_gelu(ctx->ggml_ctx, h); h = conv5->forward(ctx, h); @@ -179,17 +180,18 @@ namespace MageVAE { } }; - inline ggml_tensor* replicate_pad_right_bottom(ggml_context* ctx, + inline ggml_tensor* replicate_pad_right_bottom(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int pad_w, int pad_h) { + auto ctx = runner_ctx->ggml_ctx; if (pad_w > 0) { - auto edge = ggml_ext_slice(ctx, x, 0, x->ne[0] - 1, x->ne[0]); + auto edge = ggml_ext_slice(runner_ctx, x, 0, x->ne[0] - 1, x->ne[0]); edge = ggml_repeat_4d(ctx, edge, pad_w, x->ne[1], x->ne[2], x->ne[3]); x = ggml_concat(ctx, x, edge, 0); } if (pad_h > 0) { - auto edge = ggml_ext_slice(ctx, x, 1, x->ne[1] - 1, x->ne[1]); + auto edge = ggml_ext_slice(runner_ctx, x, 1, x->ne[1] - 1, x->ne[1]); edge = ggml_repeat_4d(ctx, edge, x->ne[0], pad_h, x->ne[2], x->ne[3]); x = ggml_concat(ctx, x, edge, 1); } @@ -209,25 +211,27 @@ namespace MageVAE { blocks["proj_out"] = std::make_shared(channels, channels, std::pair{1, 1}); } - ggml_tensor* to_patches(ggml_context* ctx, ggml_tensor* x) { - x = DiT::patchify(ctx, x, patch_size, patch_size); - x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]); + ggml_tensor* to_patches(GGMLRunnerContext* runner_ctx, ggml_tensor* x) { + auto ctx = runner_ctx->ggml_ctx; + x = DiT::patchify(runner_ctx, x, patch_size, patch_size); + x = ggml_reshape_4d(ctx, x, patch_size * patch_size, channels, x->ne[1], x->ne[2]); // [N, np, C, P] -> [N, np, P, C] for attention over P pixels. - x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(runner_ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); return ggml_reshape_3d(ctx, x, channels, patch_size * patch_size, x->ne[2] * x->ne[3]); } - ggml_tensor* from_patches(ggml_context* ctx, + ggml_tensor* from_patches(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_count, int64_t batch_size, int64_t h_patches, int64_t w_patches) { - x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size); + auto ctx = runner_ctx->ggml_ctx; + x = ggml_reshape_4d(ctx, x, channels, patch_size * patch_size, patch_count, batch_size); // [N, np, P, C] -> [N, np, C, P] before spatial unpatchify. - x = ggml_ext_cont(ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(runner_ctx, ggml_permute(ctx, x, 1, 0, 2, 3)); x = ggml_reshape_3d(ctx, x, patch_size * patch_size * channels, patch_count, batch_size); - return DiT::unpatchify(ctx, x, h_patches, w_patches, patch_size, patch_size); + return DiT::unpatchify(runner_ctx, x, h_patches, w_patches, patch_size, patch_size); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* input) override { @@ -247,19 +251,19 @@ namespace MageVAE { int64_t np = wp * hp; auto h = norm->forward(ctx, input); - auto q = replicate_pad_right_bottom(ctx->ggml_ctx, q_proj->forward(ctx, h), pad_w, pad_h); - auto k = replicate_pad_right_bottom(ctx->ggml_ctx, k_proj->forward(ctx, h), pad_w, pad_h); - auto v = replicate_pad_right_bottom(ctx->ggml_ctx, v_proj->forward(ctx, h), pad_w, pad_h); - q = to_patches(ctx->ggml_ctx, q); - k = to_patches(ctx->ggml_ctx, k); - v = to_patches(ctx->ggml_ctx, v); + auto q = replicate_pad_right_bottom(ctx, q_proj->forward(ctx, h), pad_w, pad_h); + auto k = replicate_pad_right_bottom(ctx, k_proj->forward(ctx, h), pad_w, pad_h); + auto v = replicate_pad_right_bottom(ctx, v_proj->forward(ctx, h), pad_w, pad_h); + q = to_patches(ctx, q); + k = to_patches(ctx, k); + v = to_patches(ctx, v); h = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); - h = from_patches(ctx->ggml_ctx, h, np, batch, hp, wp); + h = from_patches(ctx, h, np, batch, hp, wp); if (pad_h > 0) { - h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, height); + h = ggml_ext_slice(ctx, h, 1, 0, height); } if (pad_w > 0) { - h = ggml_ext_slice(ctx->ggml_ctx, h, 0, 0, width); + h = ggml_ext_slice(ctx, h, 0, 0, width); } return ggml_add(ctx->ggml_ctx, input, proj_out->forward(ctx, h)); } @@ -311,11 +315,11 @@ namespace MageVAE { cond = std::dynamic_pointer_cast(blocks["head_blocks." + std::to_string(i)])->forward(ctx, cond); } cond = std::dynamic_pointer_cast(blocks["proj_down"])->forward(ctx, cond); - auto z = ggml_ext_zeros(ctx->ggml_ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]); + auto z = ggml_ext_zeros(ctx, cond->ne[0], cond->ne[1], LATENT_CHANNELS, cond->ne[3]); z = std::dynamic_pointer_cast(blocks["z_proj"])->forward(ctx, z); z = ggml_concat(ctx->ggml_ctx, cond, z, 2); z = std::dynamic_pointer_cast(blocks["fuse_proj"])->forward(ctx, z); - auto t = ggml_ext_zeros(ctx->ggml_ctx, image->ne[3], 1, 1, 1); + auto t = ggml_ext_zeros(ctx, image->ne[3], 1, 1, 1); t = ggml_reshape_1d(ctx->ggml_ctx, t, image->ne[3]); auto c = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); for (int i = 0; i < 21; ++i) { @@ -336,9 +340,9 @@ namespace MageVAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* condition) { auto params = std::dynamic_pointer_cast(blocks["adaLN_modulation.1"])->forward(ctx, ggml_silu(ctx->ggml_ctx, condition)); - auto chunks = ggml_ext_chunk(ctx->ggml_ctx, params, 3, 0); + auto chunks = ggml_ext_chunk(ctx, params, 3, 0); auto h = std::dynamic_pointer_cast(blocks["in_ln"])->forward(ctx, x); - h = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1))), chunks[0]); + h = ggml_add(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, h, ggml_add(ctx->ggml_ctx, chunks[1], ggml_ext_ones(ctx, 1, 1, 1, 1))), chunks[0]); h = std::dynamic_pointer_cast(blocks["mlp.0"])->forward(ctx, h); h = ggml_silu(ctx->ggml_ctx, h); h = std::dynamic_pointer_cast(blocks["mlp.2"])->forward(ctx, h); @@ -373,8 +377,8 @@ namespace MageVAE { int64_t n = cond->ne[3]; int64_t length = w * h; - auto image = ggml_ext_zeros(ctx->ggml_ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n); - auto t = ggml_ext_zeros(ctx->ggml_ctx, n, 1, 1, 1); + auto image = ggml_ext_zeros(ctx, w * PATCH_SIZE, h * PATCH_SIZE, 3, n); + auto t = ggml_ext_zeros(ctx, n, 1, 1, 1); t = ggml_reshape_1d(ctx->ggml_ctx, t, n); auto c = std::dynamic_pointer_cast(blocks["t_embedder"])->forward(ctx, t); @@ -385,15 +389,15 @@ namespace MageVAE { s = std::dynamic_pointer_cast(blocks["blocks." + std::to_string(i)])->forward(ctx, s, c); } // [N, C, H, W] -> [N*H*W, C]. - s = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3)); + s = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, s, 1, 2, 0, 3)); s = ggml_reshape_2d(ctx->ggml_ctx, s, HIDDEN_SIZE, length * n); auto y = std::dynamic_pointer_cast(blocks["y_embedder_x"])->forward(ctx, cond); // Split 32*P channels as [32, P], then produce [N*L, P, 32]. y = ggml_reshape_4d(ctx->ggml_ctx, y, length, PATCH_SIZE * PATCH_SIZE, 32, n); - y = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3)); + y = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, y, 2, 1, 0, 3)); y = ggml_reshape_3d(ctx->ggml_ctx, y, 32, PATCH_SIZE * PATCH_SIZE, length * n); - auto zeros = ggml_ext_zeros(ctx->ggml_ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1); + auto zeros = ggml_ext_zeros(ctx, 3, PATCH_SIZE * PATCH_SIZE, length * n, 1); dct = ggml_repeat_4d(ctx->ggml_ctx, dct, 64, PATCH_SIZE * PATCH_SIZE, length * n, 1); auto x = ggml_concat(ctx->ggml_ctx, zeros, y, 0); x = ggml_concat(ctx->ggml_ctx, x, dct, 0); @@ -409,9 +413,9 @@ namespace MageVAE { x = std::dynamic_pointer_cast(blocks["final_layer.linear"])->forward(ctx, x); // [N*L, P, 3] -> [N, L, 3*P] for fold/unpatchify. x = ggml_reshape_4d(ctx->ggml_ctx, x, 3, PATCH_SIZE * PATCH_SIZE, length, n); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); x = ggml_reshape_3d(ctx->ggml_ctx, x, 3 * PATCH_SIZE * PATCH_SIZE, length, n); - return DiT::unpatchify(ctx->ggml_ctx, x, h, w, PATCH_SIZE, PATCH_SIZE); + return DiT::unpatchify(ctx, x, h, w, PATCH_SIZE, PATCH_SIZE); } }; diff --git a/src/model/vae/minimax_h3_audio_vae.hpp b/src/model/vae/minimax_h3_audio_vae.hpp index 2dc8c07e2d..d5eb0d0fbf 100644 --- a/src/model/vae/minimax_h3_audio_vae.hpp +++ b/src/model/vae/minimax_h3_audio_vae.hpp @@ -28,7 +28,7 @@ namespace MiniMaxH3 { auto alpha = params["alpha"]; auto oscillation = ggml_sin(ctx->ggml_ctx, ggml_mul(ctx->ggml_ctx, x, alpha)); oscillation = ggml_mul(ctx->ggml_ctx, oscillation, oscillation); - auto eps = ggml_ext_scale(ctx->ggml_ctx, ggml_ext_ones(ctx->ggml_ctx, 1, 1, 1, 1), 1e-9f); + auto eps = ggml_ext_scale(ctx, ggml_ext_ones(ctx, 1, 1, 1, 1), 1e-9f); return ggml_add(ctx->ggml_ctx, x, ggml_div(ctx->ggml_ctx, oscillation, ggml_add(ctx->ggml_ctx, alpha, eps))); @@ -59,7 +59,7 @@ namespace MiniMaxH3 { auto h = conv2->forward(ctx, act2->forward(ctx, conv1->forward(ctx, act1->forward(ctx, x)))); if (x->ne[0] != h->ne[0]) { int64_t pad = (x->ne[0] - h->ne[0]) / 2; - x = ggml_ext_slice(ctx->ggml_ctx, x, 0, pad, x->ne[0] - pad); + x = ggml_ext_slice(ctx, x, 0, pad, x->ne[0] - pad); } return ggml_add(ctx->ggml_ctx, x, h); } @@ -162,7 +162,7 @@ namespace MiniMaxH3 { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { auto qkv_layer = std::dynamic_pointer_cast(blocks["qkv"]); auto proj = std::dynamic_pointer_cast(blocks["proj"]); - auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_layer->forward(ctx, x), 3, 0); + auto qkv = ggml_ext_chunk(ctx, qkv_layer->forward(ctx, x), 3, 0); auto bias_shape = [&](ggml_tensor* bias) { return ggml_reshape_4d(ctx->ggml_ctx, bias, bias->ne[0], 1, 1, 1); }; @@ -172,7 +172,7 @@ namespace MiniMaxH3 { int64_t sequence = x->ne[1]; auto mask = ggml_diag_mask_inf(ctx->ggml_ctx, - ggml_ext_zeros(ctx->ggml_ctx, sequence, sequence, 1, 1), + ggml_ext_zeros(ctx, sequence, sequence, 1, 1), 0); auto attn_out = ggml_ext_attention_ext(ctx, q, @@ -323,7 +323,7 @@ namespace MiniMaxH3 { auto value = block->forward(ctx, x); sum = sum == nullptr ? value : ggml_add(ctx->ggml_ctx, sum, value); } - x = ggml_ext_scale(ctx->ggml_ctx, sum, 1.f / num_kernels); + x = ggml_ext_scale(ctx, sum, 1.f / num_kernels); } auto activation = std::dynamic_pointer_cast(blocks["activation_post"]); auto conv_post = std::dynamic_pointer_cast(blocks["conv_post"]); @@ -412,7 +412,7 @@ namespace MiniMaxH3 { streams); ggml_tensor* waveform = nullptr; for (int64_t stream = 0; stream < streams; ++stream) { - auto stream_latent = ggml_ext_slice(ctx->ggml_ctx, latent, 2, stream, stream + 1); + auto stream_latent = ggml_ext_slice(ctx, latent, 2, stream, stream + 1); auto stream_waveform = decoder->forward(ctx, dec_in->forward(ctx, stream_latent)); waveform = waveform == nullptr ? stream_waveform diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index d5bb01a4db..5b0269390d 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -66,8 +66,8 @@ namespace MiniMaxH3VAE { auto reflect_pad = [&](ggml_tensor* value, int dim, int amount) { for (int i = 0; i < amount; ++i) { GGML_ASSERT(value->ne[dim] > 1); - auto left = ggml_ext_slice(ctx->ggml_ctx, value, dim, 1, 2); - auto right = ggml_ext_slice(ctx->ggml_ctx, + auto left = ggml_ext_slice(ctx, value, dim, 1, 2); + auto right = ggml_ext_slice(ctx, value, dim, value->ne[dim] - 2, @@ -82,8 +82,7 @@ namespace MiniMaxH3VAE { x = reflect_pad(x, 1, std::get<1>(temporal_padding)); int temporal_pad = std::get<0>(temporal_padding) * 2; if (temporal_pad > 0) { - x = ggml_ext_pad_ext(ctx->ggml_ctx, - ctx->backend, + x = ggml_ext_pad_ext(ctx, x, 0, 0, @@ -105,7 +104,7 @@ namespace MiniMaxH3VAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { ggml_tensor* result = nullptr; for (int64_t t = 0; t < x->ne[2]; ++t) { - auto frame = ggml_ext_slice(ctx->ggml_ctx, x, 2, t, t + 1); + auto frame = ggml_ext_slice(ctx, x, 2, t, t + 1); GGML_ASSERT(frame->ne[3] % num_channels == 0); int64_t batch_size = frame->ne[3] / num_channels; frame = ggml_cont(ctx->ggml_ctx, frame); @@ -146,9 +145,9 @@ namespace MiniMaxH3VAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { if (spatial_stride == 2) { GGML_ASSERT(x->ne[0] > 1 && x->ne[1] > 1); - auto right = ggml_ext_slice(ctx->ggml_ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1); + auto right = ggml_ext_slice(ctx, x, 0, x->ne[0] - 2, x->ne[0] - 1); x = ggml_concat(ctx->ggml_ctx, x, right, 0); - auto bottom = ggml_ext_slice(ctx->ggml_ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1); + auto bottom = ggml_ext_slice(ctx, x, 1, x->ne[1] - 2, x->ne[1] - 1); x = ggml_concat(ctx->ggml_ctx, x, bottom, 1); } return std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); @@ -263,14 +262,14 @@ namespace MiniMaxH3VAE { auto ctx = runner_ctx->ggml_ctx; int64_t rot_dim = pe->ne[2] * 2; auto rotated = Rope::apply_rope(runner_ctx, - ggml_ext_slice(ctx, x, 0, 0, rot_dim), + ggml_ext_slice(runner_ctx, x, 0, 0, rot_dim), pe, false); if (rot_dim == x->ne[0]) { return rotated; } auto tail = attention_layout(ctx, - ggml_ext_slice(ctx, x, 0, rot_dim, x->ne[0])); + ggml_ext_slice(runner_ctx, x, 0, rot_dim, x->ne[0])); return ggml_concat(ctx, rotated, tail, 0); } @@ -298,7 +297,7 @@ namespace MiniMaxH3VAE { num_head, sequence, batch_size); - auto qkv = ggml_ext_chunk(ctx->ggml_ctx, qkv_projection, 3, 0); + auto qkv = ggml_ext_chunk(ctx, qkv_projection, 3, 0); auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv[0], head_dim, @@ -434,7 +433,7 @@ namespace MiniMaxH3VAE { GGML_ASSERT(batch_size == 1); z = ggml_cont(ctx->ggml_ctx, - ggml_ext_torch_permute(ctx->ggml_ctx, z, 3, 0, 1, 2)); + ggml_ext_torch_permute(ctx, z, 3, 0, 1, 2)); z = ggml_reshape_3d(ctx->ggml_ctx, z, 24, @@ -444,8 +443,8 @@ namespace MiniMaxH3VAE { auto h = x_embedder->forward(ctx, z); int64_t num_patches = h->ne[1]; h = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1); - auto zero = ggml_ext_scale(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, 1), + auto zero = ggml_ext_scale(ctx, + ggml_ext_slice(ctx, h, 1, 0, 1), 0.f); h = ggml_concat(ctx->ggml_ctx, h, zero, 1); @@ -461,8 +460,8 @@ namespace MiniMaxH3VAE { auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); h = proj_out->forward(ctx, norm_out->forward(ctx, h)); - h = ggml_ext_slice(ctx->ggml_ctx, h, 1, 0, num_patches); - return DiT::unpatchify_3d(ctx->ggml_ctx, + h = ggml_ext_slice(ctx, h, 1, 0, num_patches); + return DiT::unpatchify_3d(ctx, h, num_frames, height, @@ -499,7 +498,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 single_frame ? moments : ggml_ext_slice(ctx->ggml_ctx, moments, 3, 0, 24); + return single_frame ? moments : ggml_ext_slice(ctx, moments, 3, 0, 24); } ggml_tensor* decode(GGMLRunnerContext* ctx, diff --git a/src/model/vae/seedvr2_vae.hpp b/src/model/vae/seedvr2_vae.hpp index 5f38aebc18..01666c355b 100644 --- a/src/model/vae/seedvr2_vae.hpp +++ b/src/model/vae/seedvr2_vae.hpp @@ -45,20 +45,20 @@ namespace SeedVR2 { // equivalent spatial filter to F16, instead of copying and convolving the halo. auto w = ggml_cast(c, params["weight"], GGML_TYPE_F32); w = ggml_reshape_3d(c, w, kw * kh, kt, in_channels * out_channels); - auto spatial = ggml_ext_slice(c, w, 1, 0, 1, false); + auto spatial = ggml_ext_slice(ctx, w, 1, 0, 1, false); for (int i = 1; i < kt; ++i) - spatial = ggml_add(c, spatial, ggml_ext_slice(c, w, 1, i, i + 1, false)); + spatial = ggml_add(c, spatial, ggml_ext_slice(ctx, w, 1, i, i + 1, false)); w = ggml_cast(c, ggml_reshape_4d(c, spatial, kw, kh, 1, in_channels * out_channels), GGML_TYPE_F16); if (ctx->conv2d_direct_enabled) { w = ggml_reshape_4d(c, w, kw, kh, in_channels, out_channels); - x = ggml_reshape_4d(c, ggml_ext_cont(c, x), x->ne[0], x->ne[1], in_channels, 1); - x = ggml_ext_conv_2d(c, x, w, bias ? params["bias"] : nullptr, + x = ggml_reshape_4d(c, ggml_ext_cont(ctx, x), x->ne[0], x->ne[1], in_channels, 1); + x = ggml_ext_conv_2d(ctx, x, w, bias ? params["bias"] : nullptr, std::get<2>(stride), std::get<1>(stride), std::get<2>(padding), std::get<1>(padding), std::get<2>(dilation), std::get<1>(dilation), true); return ggml_reshape_4d(c, x, x->ne[0], x->ne[1], 1, out_channels); } - return ggml_ext_conv_3d(c, ctx->backend, x, w, bias ? params["bias"] : nullptr, in_channels, + return ggml_ext_conv_3d(ctx, x, w, bias ? params["bias"] : nullptr, in_channels, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), std::get<2>(padding), std::get<1>(padding), 0, std::get<2>(dilation), std::get<1>(dilation), 1, force_prec_f32, true); @@ -66,8 +66,8 @@ namespace SeedVR2 { if (temporal_padding_) { const int64_t w = x->ne[0], h = x->ne[1], t = x->ne[2], ch = x->ne[3]; // Merge spatial axes so concat does not launch separately for every channel. - x = ggml_reshape_3d(c, ggml_ext_cont(c, x), w * h, t, ch); - auto head = ggml_ext_slice(c, x, 1, 0, 1); + x = ggml_reshape_3d(c, ggml_ext_cont(ctx, x), w * h, t, ch); + auto head = ggml_ext_slice(ctx, x, 1, 0, 1); auto shape = ggml_new_tensor_3d(c, x->type, w * h, temporal_padding_, ch); x = ggml_concat(c, ggml_repeat(c, head, shape), x, 1); x = ggml_reshape_4d(c, x, w, h, t + temporal_padding_, ch); @@ -92,16 +92,16 @@ namespace SeedVR2 { ~RestorePadding() { target = value; } } restore{padding, padding}; padding = {0, 0, pw}; - auto output = ggml_ext_zeros(c, ow, oh, ot, out_channels); + auto output = ggml_ext_zeros(ctx, ow, oh, ot, out_channels); for (int64_t t = 0; t < ot; ++t) { - auto frame = ggml_ext_slice(c, x, 2, t * st, t * st + kt, false); + auto frame = ggml_ext_slice(ctx, x, 2, t * st, t * st + kt, false); for (int64_t y = 0; y < oh; y += rows) { const int64_t end = std::min(oh, y + rows); const int64_t top = y * sh - ph, bottom = (end - 1) * sh - ph + kh; - auto strip = ggml_ext_slice(c, frame, 1, std::max(0, top), std::min(x->ne[1], bottom), false); - strip = ggml_ext_pad_ext(c, ctx->backend, strip, 0, 0, int(std::max(0, -top)), + auto strip = ggml_ext_slice(ctx, frame, 1, std::max(0, top), std::min(x->ne[1], bottom), false); + strip = ggml_ext_pad_ext(ctx, strip, 0, 0, int(std::max(0, -top)), int(std::max(0, bottom - x->ne[1])), 0, 0, 0, 0); - strip = ggml_ext_cont(c, strip); + strip = ggml_ext_cont(ctx, strip); strip = Conv3d::forward(ctx, strip); output = ggml_set_inplace(c, output, strip, output->nb[1], output->nb[2], output->nb[3], y * output->nb[1] + t * output->nb[2]); } @@ -153,7 +153,7 @@ namespace SeedVR2 { h = ggml_reshape_3d(c, h, width * height, frames, channels); h = ggml_cont(c, ggml_permute(c, h, 1, 2, 0, 3)); auto linear = [&](const char* name, ggml_tensor* value) { return std::dynamic_pointer_cast(blocks[name])->forward(ctx, value); }; - h = ggml_ext_attention_ext(c, ctx->backend, linear("to_q", h), linear("to_k", h), linear("to_v", h), 1, nullptr, false, ctx->flash_attn_enabled); + h = ggml_ext_attention_ext(ctx, linear("to_q", h), linear("to_k", h), linear("to_v", h), 1, nullptr, false, ctx->flash_attn_enabled); h = linear("to_out.0", h); h = ggml_cont(c, ggml_permute(c, h, 2, 0, 1, 3)); return ggml_add(c, x, ggml_reshape_4d(c, h, width, height, frames, channels)); @@ -184,8 +184,8 @@ namespace SeedVR2 { x = ggml_cont(c, ggml_permute(c, x, 0, 2, 3, 1)); x = ggml_reshape_4d(c, x, w * 2, h * 2, t * temporal_, ch); if (temporal_ == 2) { - auto first = ggml_ext_slice(c, x, 2, 0, 1); - x = t == 1 ? first : ggml_concat(c, first, ggml_ext_slice(c, x, 2, 2, t * 2), 2); + auto first = ggml_ext_slice(ctx, x, 2, 0, 1); + x = t == 1 ? first : ggml_concat(c, first, ggml_ext_slice(ctx, x, 2, 2, t * 2), 2); } return std::dynamic_pointer_cast(blocks["conv"])->forward(ctx, x); } @@ -255,7 +255,7 @@ namespace SeedVR2 { mid(); x = ggml_silu(c, causal_norm(ctx, x, blocks[root + "conv_norm_out"])); apply("conv_out"); - return decode ? x : ggml_ext_slice(c, x, 3, 0, 16); + return decode ? x : ggml_ext_slice(ctx, x, 3, 0, 16); } }; diff --git a/src/model/vae/tae.hpp b/src/model/vae/tae.hpp index a0700cab34..09f0a1c77a 100644 --- a/src/model/vae/tae.hpp +++ b/src/model/vae/tae.hpp @@ -179,9 +179,9 @@ class TinyDecoder : public UnaryBlock { // z: [n, z_channels, h, w] // return: [n, out_channels, h*8, w*8] - auto h = ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f); + auto h = ggml_ext_scale(ctx, z, 1.0f / 3.0f); h = ggml_tanh_inplace(ctx->ggml_ctx, h); - h = ggml_ext_scale(ctx->ggml_ctx, h, 3.0f); + h = ggml_ext_scale(ctx, h, 3.0f); for (int i = 0; i < num_blocks * 3 + 10; i++) { if (blocks.find(std::to_string(i)) == blocks.end()) { @@ -322,10 +322,8 @@ class WideMemBlock : public GGMLBlock { }; ggml_tensor* -patchify(ggml_context* ctx, - ggml_tensor* x, - int64_t patch_size, - int64_t b = 1) { +patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) { + auto ctx = runner_ctx->ggml_ctx; // x: [f, b*c, h*q, w*r] // return: [f, b*c*r*q, h, w] if (patch_size == 1) { @@ -342,19 +340,17 @@ patchify(ggml_context* ctx, int64_t w = W / r; int64_t h = H / q; - x = ggml_reshape_4d(ctx, x, W, q, h, C * f); // [W, q, h, C*f] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [W, h, q, C*f] - x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f); // [r, w, h, q*C*f] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [w, h, r, q*C*f] - x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f); // [f, b*c*r*q, h, w] + x = ggml_reshape_4d(ctx, x, W, q, h, C * f); // [W, q, h, C*f] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [W, h, q, C*f] + x = ggml_reshape_4d(ctx, x, r, w, h, q * C * f); // [r, w, h, q*C*f] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3)); // [w, h, r, q*C*f] + x = ggml_reshape_4d(ctx, x, w, h, r * q * C, f); // [f, b*c*r*q, h, w] return x; } -ggml_tensor* unpatchify(ggml_context* ctx, - ggml_tensor* x, - int64_t patch_size, - int64_t b = 1) { +ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) { + auto ctx = runner_ctx->ggml_ctx; // x: [f, b*c*r*q, h, w] // return: [f, b*c, h*q, w*r] if (patch_size == 1) { @@ -367,10 +363,10 @@ ggml_tensor* unpatchify(ggml_context* ctx, int64_t h = x->ne[1]; int64_t w = x->ne[0]; - x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f); // [q*c*b*f, r, h, w] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [r, w, h, q*c*b*f] - x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f); // [c*b*f, q, h, r*w] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [r*w, q, h, c*b*f] + x = ggml_reshape_4d(ctx, x, w, h, r, q * c * b * f); // [q*c*b*f, r, h, w] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3)); // [r, w, h, q*c*b*f] + x = ggml_reshape_4d(ctx, x, r * w, h, q, c * b * f); // [c*b*f, q, h, r*w] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [r*w, q, h, c*b*f] x = ggml_reshape_4d(ctx, x, r * w, q * h, c * b, f); return x; @@ -412,7 +408,7 @@ class TinyVideoEncoder : public UnaryBlock { auto first_conv = std::dynamic_pointer_cast(blocks["0"]); if (patch_size > 1) { - z = patchify(ctx->ggml_ctx, z, patch_size, 1); + z = patchify(ctx, z, patch_size, 1); } auto h = first_conv->forward(ctx, z); @@ -427,7 +423,7 @@ class TinyVideoEncoder : public UnaryBlock { h = conv->forward(ctx, h); for (int j = 0; j < num_blocks; j++) { auto block = std::dynamic_pointer_cast(blocks[std::to_string(index++)]); - auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0); + auto mem = ggml_ext_pad_ext(ctx, h, 0, 0, 0, 0, 0, 0, 1, 0); mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0); h = block->forward(ctx, h, mem); } @@ -487,9 +483,9 @@ class TinyVideoDecoder : public UnaryBlock { auto first_conv = std::dynamic_pointer_cast(blocks["1"]); // Clamp() - auto h = ggml_ext_scale(ctx->ggml_ctx, + auto h = ggml_ext_scale(ctx, ggml_tanh_inplace(ctx->ggml_ctx, - ggml_ext_scale(ctx->ggml_ctx, z, 1.0f / 3.0f)), + ggml_ext_scale(ctx, z, 1.0f / 3.0f)), 3.0f, true); @@ -498,7 +494,7 @@ class TinyVideoDecoder : public UnaryBlock { int index = 3; for (int i = 0; i < num_layers; i++) { for (int j = 0; j < num_blocks; j++) { - auto mem = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, h, 0, 0, 0, 0, 0, 0, 1, 0); + auto mem = ggml_ext_pad_ext(ctx, h, 0, 0, 0, 0, 0, 0, 1, 0); mem = ggml_view_4d(ctx->ggml_ctx, mem, h->ne[0], h->ne[1], h->ne[2], h->ne[3], h->nb[1], h->nb[2], h->nb[3], 0); if (is_wide) { auto block = std::dynamic_pointer_cast(blocks[std::to_string(index++)]); @@ -521,7 +517,7 @@ class TinyVideoDecoder : public UnaryBlock { auto last_conv = std::dynamic_pointer_cast(blocks[std::to_string(++index)]); h = last_conv->forward(ctx, h); if (patch_size > 1) { - h = unpatchify(ctx->ggml_ctx, h, patch_size, 1); + h = unpatchify(ctx, h, patch_size, 1); } // shape(W, H, 3, (t_upscale - 1) + T) => shape(W, H, 3, T) h = ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], h->ne[2], h->ne[3] - (t_upscale - 1), h->nb[1], h->nb[2], h->nb[3], (t_upscale - 1) * h->nb[3]); @@ -579,7 +575,7 @@ class TAEHV : public GGMLBlock { int64_t chunk_frames = 5 * decoder->t_upscale; int64_t pad = (chunk_frames - (num_frames % chunk_frames)) % chunk_frames; - result = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, result, 0, 0, 0, 0, 0, 0, 0, static_cast(pad), false, false); + result = ggml_ext_pad_ext(ctx, result, 0, 0, 0, 0, 0, 0, 0, static_cast(pad), false, false); int64_t num_chunks = (num_frames + pad) / chunk_frames; auto to_trim = decoder->t_upscale - 1; @@ -591,7 +587,7 @@ class TAEHV : public GGMLBlock { i * chunk_frames * result->nb[3]); to_concat.push_back(chunk); } - result = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3); + result = ggml_ext_vec_concat(ctx, to_concat, 3); result = ggml_view_4d(ctx->ggml_ctx, result, result->ne[0], result->ne[1], result->ne[2], result->ne[3] - decoder->t_upscale * 3, @@ -645,7 +641,7 @@ class TAEHV : public GGMLBlock { to_concat.push_back(chunk_padded); } - ggml_tensor* x_in = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3); + ggml_tensor* x_in = ggml_ext_vec_concat(ctx, to_concat, 3); out = encoder->forward(ctx, x_in); } else { std::vector to_concat = {}; @@ -662,7 +658,7 @@ class TAEHV : public GGMLBlock { to_concat.push_back(chunk_out); } - out = ggml_ext_vec_concat(ctx->ggml_ctx, to_concat, 3); + out = ggml_ext_vec_concat(ctx, to_concat, 3); } // Return x[:, :-3] - drop the last 3 elements in the T dimension @@ -738,21 +734,21 @@ class TAESD : public GGMLBlock { ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* z) { auto decoder = std::dynamic_pointer_cast(blocks["decoder.layers"]); if (taef2) { - z = unpatchify(ctx->ggml_ctx, z, 2); + z = unpatchify(ctx, z, 2); } auto x = decoder->forward(ctx, z); - return f16 ? unpatchify(ctx->ggml_ctx, x, 2) : x; + return f16 ? unpatchify(ctx, x, 2) : x; } ggml_tensor* encode(GGMLRunnerContext* ctx, ggml_tensor* x) { if (f16) { GGML_ASSERT(x->ne[0] % 2 == 0 && x->ne[1] % 2 == 0); - x = patchify(ctx->ggml_ctx, x, 2); + x = patchify(ctx, x, 2); } auto encoder = std::dynamic_pointer_cast(blocks["encoder.layers"]); auto z = encoder->forward(ctx, x); if (taef2) { - z = patchify(ctx->ggml_ctx, z, 2); + z = patchify(ctx, z, 2); } return z; } diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 6b94bcb0d5..52b09bd786 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -86,26 +86,12 @@ namespace WAN { lp2 -= (int)cache_x->ne[2]; } - x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); - if (w->ne[2] == 1 && x->ne[2] == 1 && x->ne[3] == in_channels) { - // One frame through a one-frame-deep kernel is a 2D conv; backends without - // im2col_3d (Metal) otherwise fall back to a much slower direct conv_3d. - if (!ggml_is_contiguous(x)) { - x = ggml_cont(ctx->ggml_ctx, x); - } - ggml_tensor* x2 = ggml_reshape_4d(ctx->ggml_ctx, x, x->ne[0], x->ne[1], in_channels, 1); - ggml_tensor* w2 = ggml_reshape_4d(ctx->ggml_ctx, w, w->ne[0], w->ne[1], in_channels, out_channels); - x2 = ggml_ext_conv_2d(ctx->ggml_ctx, x2, w2, b, - std::get<2>(stride), std::get<1>(stride), 0, 0, - std::get<2>(dilation), std::get<1>(dilation), - ctx->conv2d_direct_enabled, false, false, scale); - return ggml_reshape_4d(ctx->ggml_ctx, x2, x2->ne[0], x2->ne[1], 1, out_channels); - } - return ggml_ext_conv_3d(ctx->ggml_ctx, ctx->backend, x, w, b, in_channels, + x = ggml_ext_pad_ext(ctx, x, lp0, rp0, lp1, rp1, lp2, rp2, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + return ggml_ext_conv_3d(ctx, x, w, b, in_channels, std::get<2>(stride), std::get<1>(stride), std::get<0>(stride), 0, 0, 0, std::get<2>(dilation), std::get<1>(dilation), std::get<0>(dilation), - false, ctx->conv3d_direct_enabled, scale); + false, ctx->conv3d_direct_enabled, scale, true); } }; @@ -133,10 +119,10 @@ namespace WAN { ggml_tensor* w = params["gamma"]; w = ggml_reshape_1d(ctx->ggml_ctx, w, ggml_nelements(w)); - auto h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 3, 0, 1, 2)); // [ID, IH, IW, N*IC] + auto h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 3, 0, 1, 2)); // [ID, IH, IW, N*IC] h = ggml_rms_norm(ctx->ggml_ctx, h, 1e-12f); h = ggml_mul(ctx->ggml_ctx, h, w); - h = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, h, 1, 2, 3, 0)); + h = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, h, 1, 2, 3, 0)); return h; } @@ -223,16 +209,16 @@ namespace WAN { } else { auto time_conv = std::dynamic_pointer_cast(blocks["time_conv"]); - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // chunk_idx >= 2 // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } if (chunk_idx == 1 && cache_x->ne[2] < 2) { // Rep - cache_x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0); + cache_x = ggml_ext_pad_ext(ctx, cache_x, 0, 0, 0, 0, (int)cache_x->ne[2], 0, 0, 0); // aka cache_x = torch.cat([torch.zeros_like(cache_x).to(cache_x.device),cache_x],dim=2) } if (chunk_idx == 1) { @@ -241,9 +227,9 @@ namespace WAN { x = time_conv->forward(ctx, x, feat_cache[idx]); } feat_cache[idx] = cache_x; - x = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2); // (2, c, t, h*w) - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 3, 1, 2)); // (c, t, 2, h*w) - x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c); // (c, t*2, h, w) + x = ggml_reshape_4d(ctx->ggml_ctx, x, w * h, t, c, 2); // (2, c, t, h*w) + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 3, 1, 2)); // (c, t, 2, h*w) + x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, 2 * t, c); // (c, t*2, h, w) } } } @@ -252,18 +238,18 @@ namespace WAN { if (mode != "none") { auto resample_1 = std::dynamic_pointer_cast(blocks["resample.1"]); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // (t, c, h, w) + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // (t, c, h, w) if (mode == "upsample2d") { x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST); } else if (mode == "upsample3d") { x = ggml_upscale(ctx->ggml_ctx, x, 2, GGML_SCALE_MODE_NEAREST); } else if (mode == "downsample2d") { - x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); } else if (mode == "downsample3d") { - x = ggml_ext_pad(ctx->ggml_ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); + x = ggml_ext_pad(ctx, x, 1, 1, 0, 0, ctx->circular_x_enabled, ctx->circular_y_enabled); } x = resample_1->forward(ctx, x); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // (c, t, h, w) + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // (c, t, h, w) } if (mode == "downsample3d") { @@ -275,9 +261,9 @@ namespace WAN { } else { auto time_conv = std::dynamic_pointer_cast(blocks["time_conv"]); - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -1, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -1, x->ne[2]); x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), x, 2); x = time_conv->forward(ctx, x); @@ -320,19 +306,19 @@ namespace WAN { int pad_t = (factor_t - T % factor_t) % factor_t; - x = ggml_ext_pad_ext(ctx->ggml_ctx, ctx->backend, x, 0, 0, 0, 0, pad_t, 0, 0, 0); + x = ggml_ext_pad_ext(ctx, x, 0, 0, 0, 0, pad_t, 0, 0, 0); T = x->ne[2]; x = ggml_reshape_4d(ctx->ggml_ctx, x, W * H, factor_t, T / factor_t, C); // [C, T/factor_t, factor_t, H*W] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C, factor_t, T/factor_t, H*W] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [C, factor_t, T/factor_t, H*W] x = ggml_reshape_4d(ctx->ggml_ctx, x, W, factor_s, (H / factor_s) * (T / factor_t), factor_t * C); // [C*factor_t, T/factor_t*H/factor_s, factor_s, W] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C*factor_t, factor_s, T/factor_t*H/factor_s, W] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [C*factor_t, factor_s, T/factor_t*H/factor_s, W] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s, W / factor_s, (H / factor_s) * (T / factor_t), factor_s * factor_t * C); // [C*factor_t*factor_s, T/factor_t*H/factor_s, W/factor_s, factor_s] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); // [C*factor_t*factor_s, factor_s, T/factor_t*H/factor_s, W/factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [C*factor_t*factor_s, factor_s, T/factor_t*H/factor_s, W/factor_s] x = ggml_reshape_3d(ctx->ggml_ctx, x, (W / factor_s) * (H / factor_s) * (T / factor_t), group_size, out_channels); // [out_channels, group_size, T/factor_t*H/factor_s*W/factor_s] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [out_channels, T/factor_t*H/factor_s*W/factor_s, group_size] - x = ggml_mean(ctx->ggml_ctx, x); // [out_channels, T/factor_t*H/factor_s*W/factor_s, 1] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 1, 0, 2, 3)); // [out_channels, T/factor_t*H/factor_s*W/factor_s, group_size] + x = ggml_mean(ctx->ggml_ctx, x); // [out_channels, T/factor_t*H/factor_s*W/factor_s, 1] x = ggml_reshape_4d(ctx->ggml_ctx, x, W / factor_s, H / factor_s, T / factor_t, out_channels); return x; } @@ -374,15 +360,15 @@ namespace WAN { C = out_channels; x = ggml_reshape_4d(ctx->ggml_ctx, x, W, H * T, factor_s, factor_s * factor_t * C); // [C*factor_t*factor_s, factor_s, T*H, W] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [C*factor_t*factor_s, T*H, W, factor_s] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, H * T, factor_s, factor_t * C); // [C*factor_t, factor_s, T*H, W*factor_s] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [C*factor_t, T*H, factor_s, W*factor_s] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W * factor_s * H, T, factor_t, C); // [C, factor_t, T, H*factor_s*W*factor_s] - x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s] + x = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [C, T, factor_t, H*factor_s*W*factor_s] x = ggml_reshape_4d(ctx->ggml_ctx, x, factor_s * W, factor_s * H, factor_t * T, C); // [C, T*factor_t, H*factor_s, W*factor_s] if (first_chunk) { - x = ggml_ext_slice(ctx->ggml_ctx, x, 2, factor_t - 1, x->ne[2]); + x = ggml_ext_slice(ctx, x, 2, factor_t - 1, x->ne[2]); } return x; @@ -451,11 +437,11 @@ namespace WAN { if (feat_cache.size() > 0) { int idx = feat_idx; - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } @@ -629,7 +615,7 @@ namespace WAN { x = norm->forward(ctx, x); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // (t, c, h, w) + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // (t, c, h, w) const int64_t n = x->ne[3]; const int64_t c = x->ne[2]; @@ -637,28 +623,28 @@ namespace WAN { const int64_t w = x->ne[0]; auto qkv = to_qkv->forward(ctx, x); - auto qkv_vec = split_image_qkv(ctx->ggml_ctx, qkv); + auto qkv_vec = ggml_ext_split_image_qkv(ctx, qkv); auto q = qkv_vec[0]; - q = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, q, 2, 0, 1, 3)); // [t, h, w, c] - q = ggml_reshape_3d(ctx->ggml_ctx, q, c, h * w, n); // [t, h * w, c] + q = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 2, 0, 1, 3)); // [t, h, w, c] + q = ggml_reshape_3d(ctx->ggml_ctx, q, c, h * w, n); // [t, h * w, c] auto k = qkv_vec[1]; - k = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, k, 2, 0, 1, 3)); // [t, h, w, c] - k = ggml_reshape_3d(ctx->ggml_ctx, k, c, h * w, n); // [t, h * w, c] + k = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, k, 2, 0, 1, 3)); // [t, h, w, c] + k = ggml_reshape_3d(ctx->ggml_ctx, k, c, h * w, n); // [t, h * w, c] auto v = qkv_vec[2]; v = ggml_reshape_3d(ctx->ggml_ctx, v, h * w, c, n); // [t, c, h * w] - v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, v, 1, 0, 2, 3)); // [t, h * w, c] + v = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, v, 1, 0, 2, 3)); // [t, h * w, c] x = ggml_ext_attention_ext(ctx, q, k, v, 1, nullptr, false, ctx->flash_attn_enabled); // [t, h * w, c] - x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [t, c, h * w] - x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n); // [t, c, h, w] + x = ggml_ext_cont(ctx, ggml_permute(ctx->ggml_ctx, x, 1, 0, 2, 3)); // [t, c, h * w] + x = ggml_reshape_4d(ctx->ggml_ctx, x, w, h, c, n); // [t, c, h, w] x = proj->forward(ctx, x); - x = ggml_ext_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx->ggml_ctx, x, 0, 1, 3, 2)); // (c, t, h, w) + x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 1, 3, 2)); // (c, t, h, w) x = ggml_add(ctx->ggml_ctx, x, identity); return x; @@ -770,11 +756,11 @@ namespace WAN { x = conv1->forward(ctx, x); } else if (feat_cache.size() > 0) { int idx = feat_idx; - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } @@ -828,11 +814,11 @@ namespace WAN { x = head_2->forward(ctx, x); } else if (feat_cache.size() > 0) { int idx = feat_idx; - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } @@ -962,11 +948,11 @@ namespace WAN { x = conv1->forward(ctx, x); } else if (feat_cache.size() > 0) { int idx = feat_idx; - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } @@ -1020,11 +1006,11 @@ namespace WAN { x = head_2->forward(ctx, x); } else if (feat_cache.size() > 0) { int idx = feat_idx; - auto cache_x = ggml_ext_slice(ctx->ggml_ctx, x, 2, -CACHE_T, x->ne[2]); + auto cache_x = ggml_ext_slice(ctx, x, 2, -CACHE_T, x->ne[2]); if (cache_x->ne[2] < 2 && feat_cache[idx] != nullptr) { // cache last frame of last two chunk cache_x = ggml_concat(ctx->ggml_ctx, - ggml_ext_slice(ctx->ggml_ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), + ggml_ext_slice(ctx, feat_cache[idx], 2, -1, feat_cache[idx]->ne[2]), cache_x, 2); } @@ -1137,10 +1123,8 @@ namespace WAN { } } - static ggml_tensor* patchify(ggml_context* ctx, - ggml_tensor* x, - int64_t patch_size, - int64_t b = 1) { + static ggml_tensor* patchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) { + auto ctx = runner_ctx->ggml_ctx; // x: [b*c, f, h*q, w*r] // return: [b*c*r*q, f, h, w] if (patch_size == 1) { @@ -1153,21 +1137,22 @@ namespace WAN { int64_t h = x->ne[1] / q; int64_t w = x->ne[0] / r; - x = ggml_reshape_4d(ctx, x, r * w, q, h, f * c * b); // [b*c*f, h, q, w*r] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [b*c*f, q, h, w*r] - x = ggml_reshape_4d(ctx, x, r, w, h * q, f * c * b); // [b*c*f, q*h, w, r] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 1, 2, 0, 3)); // [b*c*f, r, q*h, w] - x = ggml_reshape_4d(ctx, x, w * h, q * r, f, c * b); // [b*c, f, r*q, h*w] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [b*c, r*q, f, h*w] - x = ggml_reshape_4d(ctx, x, w, h, f, q * r * c * b); // [b*c*r*q, f, h, w] + x = ggml_reshape_4d(ctx, x, r * w, q, h, f * c * b); // [b*c*f, h, q, w*r] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [b*c*f, q, h, w*r] + x = ggml_reshape_4d(ctx, x, r, w, h * q, f * c * b); // [b*c*f, q*h, w, r] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 1, 2, 0, 3)); // [b*c*f, r, q*h, w] + x = ggml_reshape_4d(ctx, x, w * h, q * r, f, c * b); // [b*c, f, r*q, h*w] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [b*c, r*q, f, h*w] + x = ggml_reshape_4d(ctx, x, w, h, f, q * r * c * b); // [b*c*r*q, f, h, w] return x; } - static ggml_tensor* unpatchify(ggml_context* ctx, + static ggml_tensor* unpatchify(GGMLRunnerContext* runner_ctx, ggml_tensor* x, int64_t patch_size, int64_t b = 1) { + auto ctx = runner_ctx->ggml_ctx; // x: [b*c*r*q, f, h, w] // return: [b*c, f, h*q, w*r] if (patch_size == 1) { @@ -1180,13 +1165,13 @@ namespace WAN { int64_t h = x->ne[1]; int64_t w = x->ne[0]; - x = ggml_reshape_4d(ctx, x, w * h, f, q * r, c * b); // [b*c, r*q, f, h*w] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [b*c, f, r*q, h*w] - x = ggml_reshape_4d(ctx, x, w, h * q, r, f * c * b); // [b*c*f, r, q*h, w] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 2, 0, 1, 3)); // [b*c*f, q*h, w, r] - x = ggml_reshape_4d(ctx, x, r * w, h, q, f * c * b); // [b*c*f, q, h, w*r] - x = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, x, 0, 2, 1, 3)); // [b*c*f, h, q, w*r] - x = ggml_reshape_4d(ctx, x, r * w, q * h, f, c * b); // [b*c, f, h*q, w*r] + x = ggml_reshape_4d(ctx, x, w * h, f, q * r, c * b); // [b*c, r*q, f, h*w] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [b*c, f, r*q, h*w] + x = ggml_reshape_4d(ctx, x, w, h * q, r, f * c * b); // [b*c*f, r, q*h, w] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 2, 0, 1, 3)); // [b*c*f, q*h, w, r] + x = ggml_reshape_4d(ctx, x, r * w, h, q, f * c * b); // [b*c*f, q, h, w*r] + x = ggml_ext_cont(runner_ctx, ggml_ext_torch_permute(runner_ctx, x, 0, 2, 1, 3)); // [b*c*f, h, q, w*r] + x = ggml_reshape_4d(ctx, x, r * w, q * h, f, c * b); // [b*c, f, h*q, w*r] return x; } @@ -1203,7 +1188,7 @@ namespace WAN { clear_cache(); - x = patchify(ctx->ggml_ctx, x, patch_size, b); + x = patchify(ctx, x, patch_size, b); // sd::ggml_graph_cut::mark_graph_cut(x, "wan_vae.encode.prelude", "x"); auto encoder = std::dynamic_pointer_cast(blocks["encoder"]); @@ -1215,11 +1200,11 @@ namespace WAN { for (int i = 0; i < iter_; i++) { _enc_conv_idx = 0; if (i == 0) { - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 0, 1); // [b*c, 1, h, w] + auto in = ggml_ext_slice(ctx, x, 2, 0, 1); // [b*c, 1, h, w] out = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i); } else { // if is_2D, drop 3 out of 4 frames - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i); // [b*c, 4, h, w] + auto in = ggml_ext_slice(ctx, x, 2, 1 + 4 * (i - 1), (is_2D ? 1 - 3 : 1) + 4 * i); // [b*c, 4, h, w] auto out_ = encoder->forward(ctx, in, b, _enc_feat_map, _enc_conv_idx, i); out = ggml_concat(ctx->ggml_ctx, out, out_, 2); } @@ -1230,7 +1215,7 @@ namespace WAN { } else { out = conv1->forward(ctx, out); } - auto mu = ggml_ext_chunk(ctx->ggml_ctx, out, 2, 3)[0]; + auto mu = ggml_ext_chunk(ctx, out, 2, 3)[0]; // sd::ggml_graph_cut::mark_graph_cut(mu, "wan_vae.encode.final", "mu"); clear_cache(); return mu; @@ -1264,10 +1249,10 @@ namespace WAN { for (int i = 0; i < iter_; i++) { _conv_idx = 0; if (i == 0) { - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w] + auto in = ggml_ext_slice(ctx, x, 2, i, i + 1); // [b*c, 1, h, w] out = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i); } else { - auto in = ggml_ext_slice(ctx->ggml_ctx, x, 2, i, i + 1); // [b*c, 1, h, w] + auto in = ggml_ext_slice(ctx, x, 2, i, i + 1); // [b*c, 1, h, w] auto out_ = decoder->forward(ctx, in, b, _feat_map, _conv_idx, i); out = ggml_concat(ctx->ggml_ctx, out, out_, 2); if (is_2D) { @@ -1278,7 +1263,7 @@ namespace WAN { } } } - out = unpatchify(ctx->ggml_ctx, out, patch_size, b); + out = unpatchify(ctx, out, patch_size, b); // sd::ggml_graph_cut::mark_graph_cut(out, "wan_vae.decode.final", "out"); clear_cache(); return out; @@ -1309,7 +1294,7 @@ namespace WAN { const int global_frame = chunk_idx + static_cast(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); + auto in = ggml_ext_slice(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) { @@ -1322,7 +1307,7 @@ namespace WAN { 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); + return unpatchify(ctx, out, patch_size, b); } }; diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index e8afbd19a7..b49ef28c4c 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -856,11 +856,13 @@ bool StableDiffusionGGML::set_sage_attention_enabled(bool enabled) { if (ctx == nullptr) { return false; } - auto* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); - auto* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); - auto* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1); - auto* op = ggml_sage_attn(ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO); - bool supported = true; + auto* q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); + auto* k = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 128, 128, 1, 1); + auto* v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 128, 128, 1, 1); + GGMLRunnerContext probe_ctx; + probe_ctx.ggml_ctx = ctx; + auto* op = ggml_ext_sage_attn(&probe_ctx, q, k, v, 1.f / sqrtf(128.f), GGML_SAGE_ATTN_AUTO); + bool supported = true; for (auto backend : backend_manager.runtime_backends(SDBackendModule::DIFFUSION)) { if (!ggml_backend_supports_op(backend, op)) { LOG_ERROR("SageAttention is unavailable on %s; it requires patched GGML, CUDA Toolkit 12.0 or newer, and SM80 or newer kernels", @@ -2442,9 +2444,9 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr Date: Tue, 10 Nov 2026 21:40:13 +0800 Subject: [PATCH 49/57] feat(server): make async job queue limits configurable (#2316) --- examples/server/README.md | 10 ++++++++++ examples/server/main.cpp | 4 ++++ examples/server/runtime.cpp | 18 ++++++++++++++++++ examples/server/runtime.h | 6 ++++-- 4 files changed, 36 insertions(+), 2 deletions(-) diff --git a/examples/server/README.md b/examples/server/README.md index d93d829bc7..33b256eb61 100644 --- a/examples/server/README.md +++ b/examples/server/README.md @@ -287,6 +287,16 @@ If you want to use a different host or port, pass: --listen-ip --listen-port ``` +The native async job queue can be tuned with: + +```bash +--max-pending-jobs --completed-ttl-seconds +``` + +`--max-pending-jobs` defaults to 64 and limits how many async jobs may be +queued or generating. `--completed-ttl-seconds` defaults to 600 and controls +how long completed jobs remain available through the async API. + # Frontend ## Build with Frontend diff --git a/examples/server/main.cpp b/examples/server/main.cpp index a10cb66f1c..67d557a4e5 100644 --- a/examples/server/main.cpp +++ b/examples/server/main.cpp @@ -113,6 +113,10 @@ int main(int argc, const char** argv) { std::vector upscaler_cache; std::mutex upscaler_mutex; AsyncJobManager async_job_manager; + async_job_manager.max_pending_jobs = + static_cast(svr_params.max_pending_jobs); + async_job_manager.completed_ttl_seconds = + static_cast(svr_params.completed_ttl_seconds); ServerRuntime runtime = { models.context.get(), &sd_ctx_mutex, diff --git a/examples/server/runtime.cpp b/examples/server/runtime.cpp index b187ff85a3..5d08bba305 100644 --- a/examples/server/runtime.cpp +++ b/examples/server/runtime.cpp @@ -351,6 +351,12 @@ ArgOptions SDSvrParams::get_options() { options.int_options = { {"", "--listen-port", "server listen port (default: 1234)", &listen_port}, + {"", "--max-pending-jobs", + "maximum number of pending async jobs (default: 64)", + &max_pending_jobs}, + {"", "--completed-ttl-seconds", + "retention time for completed async jobs in seconds (default: 600)", + &completed_ttl_seconds}, }; options.bool_options = { @@ -388,6 +394,16 @@ bool SDSvrParams::validate() { return false; } + if (max_pending_jobs <= 0) { + LOG_ERROR("error: max_pending_jobs must be greater than 0"); + return false; + } + + if (completed_ttl_seconds < 0) { + LOG_ERROR("error: completed_ttl_seconds must be non-negative"); + return false; + } + if (!serve_html_path.empty() && !fs::exists(fs::u8path(serve_html_path))) { LOG_ERROR("error: serve_html_path file does not exist: %s", serve_html_path.c_str()); return false; @@ -408,6 +424,8 @@ std::string SDSvrParams::to_string() const { << " log_level: " << log_level_name(log_level) << ",\n" << " listen_ip: " << listen_ip << ",\n" << " listen_port: \"" << listen_port << "\",\n" + << " max_pending_jobs: " << max_pending_jobs << ",\n" + << " completed_ttl_seconds: " << completed_ttl_seconds << ",\n" << " serve_html_path: \"" << serve_html_path << "\",\n" << "}"; return oss.str(); diff --git a/examples/server/runtime.h b/examples/server/runtime.h index 8986f38270..ba9ddc416b 100644 --- a/examples/server/runtime.h +++ b/examples/server/runtime.h @@ -20,8 +20,10 @@ struct AsyncJobManager; struct ServerModels; struct SDSvrParams { - std::string listen_ip = "127.0.0.1"; - int listen_port = 1234; + std::string listen_ip = "127.0.0.1"; + int listen_port = 1234; + int max_pending_jobs = 64; + int completed_ttl_seconds = 600; std::string serve_html_path; std::vector model_dirs; bool normal_exit = false; From e1ca5e8bc7f1b5b8e2cbbbefc7fdaad1caffbd32 Mon Sep 17 00:00:00 2001 From: losewayy <1134285612@qq.com> Date: Tue, 10 Nov 2026 23:08:05 +0800 Subject: [PATCH 50/57] perf(cuda): fuse bias into INT8 CUTLASS epilogues (#2320) --- docs/int8_convrot.md | 2 +- src/core/cuda/i8_gemm.cu | 50 +++++++++++++++++++++++++++++----------- src/core/cuda/i8_gemm.h | 2 +- src/core/cuda/int8.cu | 6 ++--- 4 files changed, 42 insertions(+), 18 deletions(-) diff --git a/docs/int8_convrot.md b/docs/int8_convrot.md index c849dda14e..cf6219edd4 100644 --- a/docs/int8_convrot.md +++ b/docs/int8_convrot.md @@ -94,7 +94,7 @@ The packed runtime activation tensor contains the I8 activation rows and their f - CPU provides the portable regular Hadamard, activation quantization, INT8 matrix multiplication, and scale restoration implementations. - Metal supports INT8 tensorwise and convrot on Apple silicon (Apple GPU family 7 or newer). It keeps weights in INT8 and rotates and quantizes activations on the GPU. On macOS/iOS 26 or newer, it probes Metal TensorOps `matmul2d` with INT8 inputs and INT32 cooperative results. It combines 1024-product partial sums using wrapping integer addition to preserve exact accumulation on older GPUs, then applies scales and bias in registers. Older systems or unavailable TensorOps pipelines use SIMD-group matrix multiplication: INT8 tiles are converted to FP16, and FP32 partial sums cover at most 256 products before being combined as integers, preserving the I32 accumulation contract. Convrot groups 1, 4, 16, 64, and 256 are supported. These kernels require the patched GGML Metal extension services and are enabled by `SD_METAL=ON`. -- NVIDIA CUDA devices with compute capability 7.5 or newer use the native accelerated path. For H256, CUDA fuses the rotation, row-wise maximum reduction, and activation quantization. Eligible shapes use the CUTLASS INT8 implementation when enabled; other supported shapes use cuBLAS for I8 x I8 to I32 GEMM and a CUDA kernel for scale restoration and bias addition. +- 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. For linear layers, it restores scales and adds any aligned bias before writing the result; 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. - 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. diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu index 9d06c91581..7c2702d479 100644 --- a/src/core/cuda/i8_gemm.cu +++ b/src/core/cuda/i8_gemm.cu @@ -176,9 +176,22 @@ namespace sd_i8_cutlass { dst[i] = up ? (i % 2 ? up[i / 2] : src[i / 2]) : src[i / 2 + (i % 2) * (n / 2)]; } + template + struct AddBias { + CUTLASS_DEVICE Fragment operator()(const Fragment& value, const Fragment& bias) const { + Fragment result; + // Prevent contraction with the preceding scale multiplication. + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < Fragment::kElements; ++i) + result[i] = __fadd_rn(value[i], bias[i]); + return result; + } + }; + template struct DenseTypes { static constexpr bool kSplit = false; + static constexpr bool kBias = !PackedW4; using Weight = std::conditional_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; @@ -190,7 +203,9 @@ namespace sd_i8_cutlass { 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 Add = cutlass::epilogue::threadblock::VisitorCompute; + using Biased = cutlass::epilogue::threadblock::Sm80EVT; + using Epilogue = cutlass::epilogue::threadblock::Sm80EVT>; using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor< int8_t, cutlass::layout::RowMajor, @@ -219,6 +234,7 @@ namespace sd_i8_cutlass { template struct SwigluTypes { static constexpr bool kSplit = Split; + static constexpr bool kBias = false; using Weight = int8_t; using Tile = typename Config::Tile; using Warp = typename Config::Warp; @@ -259,24 +275,30 @@ namespace sd_i8_cutlass { bool packed_w4; bool split; size_t kernel_bytes = 0, scratch_offset = 0; - int occupancy = 0; - ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*) = nullptr; + int occupancy = 0; + ggml_status (*launch)(Plan&, const int8_t*, const int8_t*, const float*, const float*, float*, void*, cudaStream_t, const int8_t*, const float*, const float*) = 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}}}; + 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, const float* bias) { + typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}}; + typename Types::Epilogue::Arguments epilogue; + if constexpr (Types::kBias) { + // A negative zero default preserves signed-zero outputs. + epilogue = {{scales, {bias, -0.f, {}}, {}}, {out, {p.n, cute::_1{}, p.rows * p.n}}}; + } else { + epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}}; + } return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, reinterpret_cast(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, const int8_t* up, const float* up_scales) { + static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const int8_t* up, const float* up_scales, const float* bias) { using Kernel = typename Types::Kernel; if ((reinterpret_cast(x) | reinterpret_cast(w) | reinterpret_cast(xs) | - reinterpret_cast(ws) | reinterpret_cast(out) | reinterpret_cast(up) | reinterpret_cast(up_scales)) % + reinterpret_cast(ws) | reinterpret_cast(out) | reinterpret_cast(up) | reinterpret_cast(up_scales) | reinterpret_cast(bias)) % 16) return GGML_STATUS_FAILED; auto* paired = static_cast(scratch); - auto args = arguments(p, x, w, xs, p.swiglu ? paired : ws, out); + auto args = arguments(p, x, w, xs, p.swiglu ? paired : ws, out, bias); if (Kernel::can_implement(args) != cutlass::Status::kSuccess) return GGML_STATUS_FAILED; if (!p.occupancy) { @@ -312,7 +334,7 @@ namespace sd_i8_cutlass { using Kernel = typename Types::Kernel; if (sizeof(typename Kernel::SharedStorage) > p.device.shared) return false; - auto args = arguments(p, nullptr, nullptr, nullptr, nullptr, nullptr); + auto args = arguments(p, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); if (Kernel::can_implement(args) != cutlass::Status::kSuccess) return false; size_t bytes = 0; @@ -394,11 +416,13 @@ SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64 return nullptr; return prepare_gemm(device, rows, 2 * n, k, true, workspace, false, true); } -ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) { - return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr); +ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream, const float* bias) { + if (bias && (plan->swiglu || plan->packed_w4)) + return GGML_STATUS_FAILED; + return plan->launch(*plan, x, w, xs, ws, out, scratch, stream, nullptr, nullptr, bias); } ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream) { - return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales); + return plan->launch(*plan, x, gate, xs, gate_scales, out, scratch, stream, up, up_scales, nullptr); } 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 index dc0d15d6ea..756238ef11 100644 --- a/src/core/cuda/i8_gemm.h +++ b/src/core/cuda/i8_gemm.h @@ -4,7 +4,7 @@ #include "ggml-ext.h" struct SDI8CudaGemmPlan; SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace, bool packed_w4 = false); -ggml_status sd_cuda_i8_gemm_execute(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream); +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, const float* bias = nullptr); // n is the width of each branch; gate/up weights and scales remain separate. SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare_split(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, size_t* workspace); ggml_status sd_cuda_i8_gemm_execute_split(SDI8CudaGemmPlan* plan, const int8_t* x, const int8_t* gate, const int8_t* up, const float* xs, const float* gate_scales, const float* up_scales, float* out, void* scratch, cudaStream_t stream); diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu index 01a614731e..8e92468482 100644 --- a/src/core/cuda/int8.cu +++ b/src/core/cuda/int8.cu @@ -447,7 +447,7 @@ namespace SD_CUDA_NAMESPACE { } #ifdef SD_CUDA_I8_CUTLASS static bool cutlass_compatible(const SDI8Operation& op, const ggml_tensor* dst) { - return op.weight_scale && !op.bias && op.weight->ne[0] <= INT_MAX / 16384 && + return op.weight_scale && op.weight->ne[0] <= INT_MAX / 16384 && aligned(op.bias, 16) && 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))); } @@ -531,7 +531,7 @@ namespace SD_CUDA_NAMESPACE { #ifdef SD_CUDA_I8_CUTLASS size_t gemm_bytes = 0; // Retain cuBLAS when I32 accumulation can overflow. - if (cutlass_compatible(op, dst) && + if (cutlass_compatible(op, dst) && (!op.bias || p->kind == MAT) && (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); @@ -819,7 +819,7 @@ namespace SD_CUDA_NAMESPACE { #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_at(p.gemm_offset), stream); + intermediate ? reinterpret_cast(accum) : gemm_output, scratch_at(p.gemm_offset), stream, bias); if (status != GGML_STATUS_SUCCESS) return status; } else From 0bb132831f4cabf988b8854d6daf80820695aeaf Mon Sep 17 00:00:00 2001 From: losewayy <1134285612@qq.com> Date: Tue, 10 Nov 2026 23:29:17 +0800 Subject: [PATCH 51/57] feat: support Sol-H3-VAE decoding for MiniMax-H3 (#2317) --- docs/minimax_h3.md | 18 ++ docs/performance.md | 10 +- src/model/vae/minimax_h3_sol_attention.hpp | 204 ++++++++++++++ src/model/vae/minimax_h3_vae.hpp | 300 ++++++++++++++++----- src/model/vae/vae_tiling.hpp | 25 +- 5 files changed, 485 insertions(+), 72 deletions(-) create mode 100644 src/model/vae/minimax_h3_sol_attention.hpp diff --git a/docs/minimax_h3.md b/docs/minimax_h3.md index b5514db69d..6b672ba99f 100644 --- a/docs/minimax_h3.md +++ b/docs/minimax_h3.md @@ -99,6 +99,24 @@ 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. +## Sol-H3-VAE decoding + +Pass a [Sol-H3-VAE checkpoint](https://huggingface.co/Comfy-Org/MiniMax-H3/tree/main/vae) +with `--vae` to use the distilled video decoder. The ComfyUI FP16 and INT8 +ConvRot checkpoints include the MiniMax-H3 encoder for image and video inputs. +Supply `--audio-vae` as usual to decode the soundtrack. + +Sol decodes the whole clip by default. Use `--vae-tiling` for spatial tiling or +`--temporal-tiling` for temporal tiling. Either can reduce memory use but may +affect output quality because the decoder has less context at tile boundaries. + +Temporal tiles default to 10 latent frames with 5 latent frames of overlap. +Set `temporal_tile_frames` and `temporal_tile_overlap` in +`--extra-tiling-args` to override them. Both must be multiples of 5 latent +frames, and overlap must be smaller than the tile. The last tile may be shorter. +Circular decoding requires spatial tiling on each circular axis, with a tile +size smaller than the input on that axis. + ## Latent hires second pass (two-stage sampling) MiniMax-H3 supports `--hires` with two interchangeable second-stage upscalers: diff --git a/docs/performance.md b/docs/performance.md index f347dc1721..1f5771c849 100644 --- a/docs/performance.md +++ b/docs/performance.md @@ -92,10 +92,12 @@ dimensions down to the minimum size. Each spatial retry must reduce the effectiv tile size. These runtime adjustments do not change the caller's parameters. Execution failures are not retried, and encoding has no automatic OOM retry. -`--temporal-tiling` remains independent of spatial tiling. MiniMax H3 always uses -spatial tiling (256x256 pixels and 25% overlap by default) and its own temporal -windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap`; explicit -spatial sizes are honored. +`--temporal-tiling` remains independent of spatial tiling. The original MiniMax-H3 +video VAE uses spatial tiling (256x256 pixels and 25% overlap by default) and its own +temporal windows. With `--vae-tiling`, its overlap follows `--vae-tile-overlap` and +explicit spatial sizes are honored. [Sol-H3-VAE](minimax_h3.md#sol-h3-vae-decoding) +decodes the whole clip by default and supports opt-in spatial and block-aligned +temporal tiling. Tiling may affect output quality. ## Offload weights to the CPU to save VRAM without reducing generation speed. diff --git a/src/model/vae/minimax_h3_sol_attention.hpp b/src/model/vae/minimax_h3_sol_attention.hpp new file mode 100644 index 0000000000..66d5b5c4fd --- /dev/null +++ b/src/model/vae/minimax_h3_sol_attention.hpp @@ -0,0 +1,204 @@ +#ifndef __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__ +#define __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__ + +#include +#include +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/tensor.hpp" + +namespace MiniMaxH3VAE { + + constexpr int64_t sol_window_size = 16; + constexpr int64_t sol_block_size = 5; + constexpr int64_t sol_global_rows = 5; + constexpr int64_t sol_spatial_shift = sol_window_size / 2; + constexpr int64_t sol_temporal_shift = 2; + + struct SolAttentionPlan { + template + struct Batch { + Tensor output_rows; + Tensor queries; + Tensor keys; + Tensor previous_queries; + Tensor previous_keys; + }; + + template + struct Block { + Tensor rows; + std::vector> batches; + }; + + std::vector>> blocks; + + SolAttentionPlan(int64_t width, int64_t height, int64_t frames, bool shifted) { + const int64_t frame_size = width * height; + const int64_t pixels = frames * frame_size; + const int64_t block_rows = sol_block_size * frame_size; + const int64_t count = (frames + sol_block_size - 1) / sol_block_size; + const int64_t shift = shifted ? sol_spatial_shift : 0; + const int64_t time_shift = shifted ? sol_temporal_shift : 0; + + auto pack = [](std::vector rows) { + const auto size = static_cast(rows.size()); + return sd::Tensor({size}, std::move(rows)); + }; + for (int64_t block = 0; block < count; ++block) { + std::vector rows; + const int64_t start = block * block_rows; + const int64_t end = std::min(pixels, start + block_rows); + for (int64_t i = start; i < end; ++i) { + rows.push_back(static_cast(i)); + } + std::vector globals; + std::vector queries; + for (int64_t i = 0; i < sol_global_rows; ++i) { + globals.push_back(static_cast(pixels + block * sol_global_rows + i)); + queries.push_back(static_cast(rows.size())); + rows.push_back(globals.back()); + } + blocks.push_back({pack(std::move(rows)), {}}); + blocks.back().batches.push_back({pack(std::move(globals)), pack(std::move(queries)), {}, {}, {}}); + } + + struct Window { + std::vector output_rows; + std::vector queries; + std::vector keys; + std::vector previous_queries; + std::vector previous_keys; + }; + for (int64_t t = -time_shift; t < frames; t += sol_block_size) { + const int64_t first = std::max(0, t); + const int64_t end = std::min(frames, t + sol_block_size); + const int64_t first_block = first / sol_block_size; + const int64_t last_block = (end - 1) / sol_block_size; + std::map> groups; + for (int64_t y = -shift; y < height; y += sol_window_size) { + for (int64_t x = -shift; x < width; x += sol_window_size) { + Window window; + for (int64_t i = first; i < end; ++i) { + for (int64_t j = std::max(0, y); j < std::min(height, y + sol_window_size); ++j) { + for (int64_t k = std::max(0, x); k < std::min(width, x + sol_window_size); ++k) { + const int64_t row = (i * height + j) * width + k; + window.output_rows.push_back(static_cast(row)); + const int64_t block = i / sol_block_size; + auto& queries = block == last_block ? window.queries : window.previous_queries; + queries.push_back(static_cast(row - block * block_rows)); + } + } + } + window.keys = window.queries; + window.previous_keys = window.previous_queries; + for (int64_t block = first_block; block <= last_block; ++block) { + const int64_t start = std::min(block_rows, pixels - block * block_rows); + auto& keys = block == last_block ? window.keys : window.previous_keys; + for (int64_t i = 0; i < sol_global_rows; ++i) { + keys.push_back(static_cast(start + i)); + } + } + groups[window.output_rows.size()].push_back(std::move(window)); + } + } + // Bound the workspace of backends that fall back to unfused attention. + constexpr size_t max_windows = 4; + for (const auto& [query_count, windows] : groups) { + for (size_t first_window = 0; first_window < windows.size(); first_window += max_windows) { + const auto n = static_cast(std::min(max_windows, windows.size() - first_window)); + auto pack_windows = [&](auto member) -> sd::Tensor { + const auto rows = static_cast((windows[first_window].*member).size()); + if (rows == 0) { + return {}; + } + sd::Tensor indices({rows, n}); + for (int64_t i = 0; i < n; ++i) { + const auto& source = windows[first_window + i].*member; + std::copy(source.begin(), source.end(), indices.data() + i * rows); + } + return indices; + }; + blocks[last_block].batches.push_back({pack_windows(&Window::output_rows), + pack_windows(&Window::queries), + pack_windows(&Window::keys), + pack_windows(&Window::previous_queries), + pack_windows(&Window::previous_keys)}); + } + } + } + } + + struct Inputs { + std::vector> blocks; + ggml_tensor* row_ids = nullptr; + + template + ggml_tensor* forward(GGMLRunnerContext* runner_ctx, + ggml_tensor* x, + ggml_tensor* pe, + int64_t heads, + Project project) const { + auto ctx = runner_ctx->ggml_ctx; + const int64_t dim = x->ne[0]; + auto output = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, x->ne[1]); + auto positions = ggml_reshape_2d(ctx, pe, pe->ne[0] * pe->ne[1] * pe->ne[2], pe->ne[3]); + auto gather = [&](ggml_tensor* rows, ggml_tensor* indices) { + auto flat = ggml_reshape_1d(ctx, indices, ggml_nelements(indices)); + return ggml_reshape_3d(ctx, ggml_get_rows(ctx, rows, flat), dim, indices->ne[0], indices->ne[1]); + }; + std::array previous{}; + for (const auto& block : blocks) { + auto local_pe = ggml_reshape_4d(ctx, ggml_get_rows(ctx, positions, block.rows), + pe->ne[0], pe->ne[1], pe->ne[2], block.rows->ne[0]); + auto current = project(ggml_get_rows(ctx, x, block.rows), local_pe); + for (const auto& batch : block.batches) { + // Gather first so a crossing window never copies either full block. + auto window = [&](size_t i, ggml_tensor* now, ggml_tensor* before) { + // Global queries read the whole native block in its existing order. + auto rows = now == nullptr ? current[i] : gather(current[i], now); + if (before != nullptr) { + rows = ggml_concat(ctx, gather(previous[i], before), rows, 1); + } + return rows; + }; + auto attended = ggml_ext_attention_ext(runner_ctx, + window(0, batch.queries, batch.previous_queries), + window(1, batch.keys, batch.previous_keys), + window(2, batch.keys, batch.previous_keys), + heads, nullptr, false, runner_ctx->flash_attn_enabled); + auto indices = ggml_reshape_1d(ctx, batch.output_rows, ggml_nelements(batch.output_rows)); + output = ggml_set_rows(ctx, output, ggml_reshape_2d(ctx, attended, dim, indices->ne[0]), indices); + } + // Later windows can only read this block and its immediate successor. + previous = current; + } + return output; + } + }; + + template + Inputs bind(BindInput bind_input) const { + Inputs inputs; + for (const auto& block : blocks) { + Block bound{bind_input(block.rows), {}}; + for (const auto& batch : block.batches) { + auto bind_optional = [&](const auto& indices) { + return indices.empty() ? nullptr : bind_input(indices); + }; + bound.batches.push_back({bind_input(batch.output_rows), bind_input(batch.queries), + bind_optional(batch.keys), bind_optional(batch.previous_queries), + bind_optional(batch.previous_keys)}); + } + inputs.blocks.push_back(std::move(bound)); + } + return inputs; + } + }; + +} // namespace MiniMaxH3VAE + +#endif // __SD_MODEL_VAE_MINIMAX_H3_SOL_ATTENTION_HPP__ diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 5b0269390d..e59eed9444 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -13,26 +13,34 @@ #include "core/rng_mt19937.hpp" #include "model/common/rope.hpp" #include "model/diffusion/dit.hpp" +#include "model/vae/minimax_h3_sol_attention.hpp" #include "model/vae/vae.hpp" namespace MiniMaxH3VAE { constexpr int H3_VIDEO_VAE_GRAPH_SIZE = 262144; + static int64_t sol_output_frames(int64_t latent_frames, int patch_size_t) { + const int64_t blocks = (latent_frames + sol_block_size - 1) / sol_block_size; + return latent_frames * patch_size_t - blocks * (patch_size_t - 1); + } + struct MiniMaxH3VAEConfig { int num_layers = 36; int patch_size_t = 4; + bool sol = false; static MiniMaxH3VAEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { MiniMaxH3VAEConfig config; - auto proj = tensors.find(prefix + ".decoder.proj_out.weight"); + config.sol = tensors.count(prefix + ".decoder.blocks.0.scale1") != 0; + 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."; + const auto layer_prefix = prefix + (config.sol ? ".decoder.blocks." : ".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); @@ -41,7 +49,12 @@ namespace MiniMaxH3VAE { 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); + } + if (config.sol || proj != tensors.end()) { + LOG_VERBOSE("minimax_h3_vae: decoder=%s, layers=%d, patch_size_t=%d", + config.sol ? "sol" : "original", + config.num_layers, + config.patch_size_t); } return config; } @@ -283,11 +296,10 @@ namespace MiniMaxH3VAE { blocks["to_out"] = std::make_shared(dim, dim, true); } - ggml_tensor* forward(GGMLRunnerContext* ctx, - ggml_tensor* x, - ggml_tensor* pe) { + std::array project(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe) { auto to_qkv = std::dynamic_pointer_cast(blocks["to_qkv"]); - auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); auto qkv_projection = to_qkv->forward(ctx, x); int64_t sequence = x->ne[1]; int64_t batch_size = x->ne[2] * x->ne[3]; @@ -320,6 +332,29 @@ namespace MiniMaxH3VAE { k = ggml_rms_norm(ctx->ggml_ctx, k, 1e-5f); q = apply_partial_rope(ctx, q, pe); k = apply_partial_rope(ctx, k, pe); + return {q, k, v}; + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, + ggml_tensor* x, + ggml_tensor* pe, + RMSNorm& norm, + const SolAttentionPlan::Inputs* sol) { + auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); + if (sol != nullptr) { + auto project_rows = [&](ggml_tensor* rows, ggml_tensor* local_pe) { + auto [q, k, v] = project(ctx, norm.forward(ctx, rows), local_pe); + auto as_rows = [&](ggml_tensor* tensor) { + return ggml_reshape_2d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, tensor), dim, rows->ne[1]); + }; + return std::array{ + as_rows(ggml_permute(ctx->ggml_ctx, q, 0, 2, 1, 3)), + as_rows(ggml_permute(ctx->ggml_ctx, k, 0, 2, 1, 3)), + as_rows(v)}; + }; + return to_out->forward(ctx, sol->forward(ctx, x, pe, num_head, project_rows)); + } + auto [q, k, v] = project(ctx, norm.forward(ctx, x), pe); auto out = ggml_ext_attention_ext(ctx, q, k, @@ -370,21 +405,33 @@ namespace MiniMaxH3VAE { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, - ggml_tensor* pe) { + ggml_tensor* pe, + const SolAttentionPlan::Inputs* sol, + int64_t rows_per_chunk) { auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); auto attn = std::dynamic_pointer_cast(blocks["attn"]); auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); auto ff = std::dynamic_pointer_cast(blocks["ff"]); - x = ggml_add(ctx->ggml_ctx, - x, - ggml_mul(ctx->ggml_ctx, - attn->forward(ctx, norm1->forward(ctx, x), pe), - params["scale1"])); - return ggml_add(ctx->ggml_ctx, - x, - ggml_mul(ctx->ggml_ctx, - ff->forward(ctx, norm2->forward(ctx, x)), - params["scale2"])); + x = ggml_add(ctx->ggml_ctx, + x, + ggml_mul(ctx->ggml_ctx, + attn->forward(ctx, x, pe, *norm1, sol), + params["scale1"])); + auto finish = [&](ggml_tensor* rows) { + auto update = ff->forward(ctx, norm2->forward(ctx, rows)); + return ggml_add(ctx->ggml_ctx, rows, ggml_mul(ctx->ggml_ctx, update, params["scale2"])); + }; + if (rows_per_chunk == 0 || x->ne[1] <= rows_per_chunk) { + return finish(x); + } + // Normalize and update each independent chunk before writing it to the layer output. + auto output = ggml_new_tensor_2d(ctx->ggml_ctx, x->type, x->ne[0], x->ne[1]); + for (int64_t row = 0; row < x->ne[1]; row += rows_per_chunk) { + auto input = ggml_ext_slice(ctx, x, 1, row, std::min(row + rows_per_chunk, x->ne[1]), false); + auto indices = ggml_ext_slice(ctx, sol->row_ids, 0, row, std::min(row + rows_per_chunk, x->ne[1]), false); + output = ggml_set_rows(ctx->ggml_ctx, output, finish(input), indices); + } + return output; } }; @@ -394,12 +441,13 @@ namespace MiniMaxH3VAE { static constexpr int num_register_tokens = 4; static constexpr int patch_size = 16; int patch_size_t; + bool sol; explicit Decoder(const MiniMaxH3VAEConfig& config) - : num_layers(config.num_layers), patch_size_t(config.patch_size_t) { + : num_layers(config.num_layers), patch_size_t(config.patch_size_t), sol(config.sol) { blocks["x_embedder"] = std::make_shared(24, dim, true); for (int i = 0; i < num_layers; ++i) { - blocks["transformer_blocks." + std::to_string(i)] = + blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)] = std::make_shared(); } blocks["norm_out"] = std::make_shared(dim, 1e-5f, true, true); @@ -418,14 +466,18 @@ namespace MiniMaxH3VAE { GGML_TYPE_F32, dim, num_register_tokens); - if (patch_size_t != 1) { + if (patch_size_t != 1 && !sol) { params["mask_token"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, dim); } } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* z, - ggml_tensor* pe) { + ggml_tensor* pe, + ggml_tensor* shifted_pe, + const std::array& sol_inputs, + ggml_tensor* pixel_mean, + ggml_tensor* pixel_std) { int64_t width = z->ne[0]; int64_t height = z->ne[1]; int64_t num_frames = z->ne[2]; @@ -442,34 +494,68 @@ namespace MiniMaxH3VAE { auto x_embedder = std::dynamic_pointer_cast(blocks["x_embedder"]); auto h = x_embedder->forward(ctx, z); int64_t num_patches = h->ne[1]; - h = ggml_concat(ctx->ggml_ctx, h, params["register_tokens"], 1); auto zero = ggml_ext_scale(ctx, ggml_ext_slice(ctx, h, 1, 0, 1), 0.f); - h = ggml_concat(ctx->ggml_ctx, h, zero, 1); + auto registers = ggml_concat(ctx->ggml_ctx, params["register_tokens"], zero, 1); + if (sol) { + const int64_t num_blocks = (num_frames + sol_block_size - 1) / sol_block_size; + registers = ggml_repeat(ctx->ggml_ctx, registers, + ggml_new_tensor_2d(ctx->ggml_ctx, registers->type, dim, num_blocks * sol_global_rows)); + } + h = ggml_concat(ctx->ggml_ctx, h, registers, 1); + const int64_t rows_per_chunk = sol ? width * height + sol_global_rows : 0; for (int i = 0; i < num_layers; ++i) { auto block = std::dynamic_pointer_cast( - blocks["transformer_blocks." + std::to_string(i)]); - h = block->forward(ctx, h, pe); + blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)]); + h = block->forward(ctx, h, sol && i % 2 != 0 ? shifted_pe : pe, sol ? &sol_inputs[i % 2] : nullptr, rows_per_chunk); sd::ggml_graph_cut::mark_graph_cut(h, "minimax_h3_vae.decoder.blocks." + std::to_string(i), "hidden_states"); } - auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); - auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); - h = proj_out->forward(ctx, norm_out->forward(ctx, h)); - h = ggml_ext_slice(ctx, h, 1, 0, num_patches); - return DiT::unpatchify_3d(ctx, - h, - num_frames, - height, - width, - patch_size_t, - patch_size, - patch_size, - true); + auto norm_out = std::dynamic_pointer_cast(blocks["norm_out"]); + auto proj_out = std::dynamic_pointer_cast(blocks["proj_out"]); + auto denormalize_and_clamp = [&](ggml_tensor* pixels) { + auto scaled = ggml_mul(ctx->ggml_ctx, pixels, pixel_std); + return ggml_clamp(ctx->ggml_ctx, ggml_add(ctx->ggml_ctx, scaled, pixel_mean), 0.f, 1.f); + }; + if (sol) { + const int64_t frame_rows = width * height; + const int64_t output_frames = sol_output_frames(num_frames, patch_size_t); + const int64_t frame_size = width * height * patch_size * patch_size; + const int64_t channels = pixel_mean->ne[3]; + auto output = ggml_new_tensor_2d(ctx->ggml_ctx, GGML_TYPE_F32, frame_size, channels * output_frames); + int64_t frame_offset = 0; + for (int64_t t = 0; t < num_frames; t += sol_block_size) { + const int64_t end = std::min(t + sol_block_size, num_frames); + auto rows = ggml_ext_slice(ctx, h, 1, t * frame_rows, end * frame_rows, false); + auto patches = proj_out->forward(ctx, norm_out->forward(ctx, rows)); + auto pixels = DiT::unpatchify_3d(ctx, patches, end - t, height, width, + patch_size_t, patch_size, patch_size, true); + // Keep only the last frame of the first latent patch in each native block. + pixels = ggml_ext_slice(ctx, pixels, 2, patch_size_t - 1, pixels->ne[2], false); + pixels = denormalize_and_clamp(pixels); + const int64_t count = pixels->ne[2]; + for (int64_t c = 0; c < channels; ++c) { + auto plane = ggml_ext_slice(ctx, pixels, 3, c, c + 1, false); + auto frames = ggml_reshape_2d(ctx->ggml_ctx, ggml_ext_cont(ctx, plane), frame_size, count); + const int64_t start = c * output_frames + frame_offset; + auto indices = ggml_ext_slice(ctx, sol_inputs[0].row_ids, 0, start, start + count, false); + output = ggml_set_rows(ctx->ggml_ctx, output, frames, indices); + } + // Write each block before expanding the next block's pixel graph. + ctx->expand_graph(output); + frame_offset += count; + } + return ggml_reshape_4d(ctx->ggml_ctx, output, width * patch_size, height * patch_size, output_frames, channels); + } + h = proj_out->forward(ctx, norm_out->forward(ctx, h)); + h = ggml_ext_slice(ctx, h, 1, 0, num_patches); + auto pixels = DiT::unpatchify_3d(ctx, h, num_frames, height, width, + patch_size_t, patch_size, patch_size, true); + return denormalize_and_clamp(pixels); } }; @@ -479,9 +565,7 @@ namespace MiniMaxH3VAE { 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, - std::tuple{1, 1, 1}); + blocks["quant_conv"] = std::make_shared(48, 48, std::tuple{1, 1, 1}); blocks["post_quant_conv"] = std::make_shared(24, 24, std::tuple{1, 1, 1}); @@ -504,15 +588,13 @@ namespace MiniMaxH3VAE { ggml_tensor* decode(GGMLRunnerContext* ctx, ggml_tensor* latent, ggml_tensor* pe, + ggml_tensor* shifted_pe, + const std::array& sol_inputs, ggml_tensor* pixel_mean, ggml_tensor* pixel_std) { auto post_quant = std::dynamic_pointer_cast(blocks["post_quant_conv"]); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); - auto pixels = decoder->forward(ctx, post_quant->forward(ctx, latent), pe); - pixels = ggml_add(ctx->ggml_ctx, - ggml_mul(ctx->ggml_ctx, pixels, pixel_std), - pixel_mean); - return ggml_clamp(ctx->ggml_ctx, pixels, 0.f, 1.f); + return decoder->forward(ctx, post_quant->forward(ctx, latent), pe, shifted_pe, sol_inputs, pixel_mean, pixel_std); } }; @@ -553,6 +635,9 @@ namespace MiniMaxH3VAE { } std::string get_desc() override { + if (config.sol) { + return "sol_h3_vae"; + } return config.patch_size_t == 1 ? "p_h3_image_vae" : "minimax_h3_video_vae"; } @@ -601,6 +686,10 @@ namespace MiniMaxH3VAE { } sd_tiling_params_t resolve_tiling_params(sd_tiling_params_t params) const override { + return config.sol ? params : resolve_original_tiling_params(params); + } + + static sd_tiling_params_t resolve_original_tiling_params(sd_tiling_params_t params) { if (!params.enabled) { params.target_overlap = 0.25f; } @@ -655,7 +744,7 @@ namespace MiniMaxH3VAE { bool circular_x = false, bool circular_y = false) override { auto input = ensure_video_shape(x); - auto tiling = resolve_tiling_params(tiling_params); + auto tiling = resolve_original_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 {}; @@ -695,6 +784,46 @@ namespace MiniMaxH3VAE { return result; } + bool supports_temporal_tiling(VAETemporalDirection direction) const override { + return config.sol && direction == VAETemporalDirection::DECODE; + } + + int get_default_temporal_tile_frames(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + return 2 * sol_block_size; + } + + int get_default_temporal_tile_overlap(VAETemporalDirection direction) const override { + SD_UNUSED(direction); + return sol_block_size; + } + + sd::Tensor _compute_temporal_tiled(int n_threads, + const sd::Tensor& input, + VAETemporalDirection direction, + const VAETemporalTilingConfig& tiling) override { + GGML_ASSERT(supports_temporal_tiling(direction)); + if (tiling.tile_frames % sol_block_size || tiling.overlap % sol_block_size || + tiling.overlap >= tiling.tile_frames) { + LOG_ERROR("Sol-H3-VAE temporal tile size and overlap must be multiples of %d latent frames, with overlap smaller than the tile", + static_cast(sol_block_size)); + return {}; + } + if (input.shape()[2] <= tiling.tile_frames) { + return _compute(n_threads, input, true); + } + auto plan = make_vae_temporal_tile_plan(input.shape()[2], tiling); + // With block-aligned starts, the same mapping gives tile lengths and output offsets. + auto frame_count = [&](int64_t frames) { return sol_output_frames(frames, config.patch_size_t); }; + return process_vae_temporal_tiles_blended(input, plan, frame_count, frame_count, + [&](const sd::Tensor& tile, const VAETemporalTile& window) { + LOG_VERBOSE("Sol-H3-VAE temporal tile %d/%d: latent frames [%lld, %lld)", + window.index + 1, static_cast(plan.tiles.size()), + static_cast(window.start), static_cast(window.end)); + return _compute(n_threads, tile, true); + }); + } + sd::Tensor decode(int n_threads, const sd::Tensor& x, sd_tiling_params_t tiling_params, @@ -702,7 +831,25 @@ namespace MiniMaxH3VAE { bool circular_x = false, bool circular_y = false, bool silent = false) override { - auto input = ensure_video_shape(x); + auto input = ensure_video_shape(x); + if (config.sol) { + circular_x = circular_x || circular_x_enabled; + circular_y = circular_y || circular_y_enabled; + if (circular_x || circular_y) { + int tile_w = static_cast(input.shape()[0]); + int tile_h = static_cast(input.shape()[1]); + float overlap; + if (tiling_params.enabled && + !get_tile_sizes(tile_w, tile_h, overlap, tiling_params, input.shape()[0], input.shape()[1])) { + return {}; + } + if ((circular_x && tile_w >= input.shape()[0]) || (circular_y && tile_h >= input.shape()[1])) { + LOG_ERROR("Sol-H3-VAE circular decoding requires spatial tiles smaller than the input on every circular axis"); + return {}; + } + } + return VAE::decode(n_threads, input, tiling_params, decode_video, circular_x, circular_y, silent); + } 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"); @@ -801,23 +948,29 @@ namespace MiniMaxH3VAE { sd::Tensor build_rope(int64_t width, int64_t height, - int64_t num_frames) { + int64_t num_frames, + bool shifted = false) { std::vector> ids; - ids.reserve(static_cast(width * height * num_frames + 5)); + const int64_t global_rows = config.sol ? ((num_frames + sol_block_size - 1) / sol_block_size) * sol_global_rows : 5; + ids.reserve(static_cast(width * height * num_frames + global_rows)); constexpr float two_pi = 6.28318530717958647692f; + const int64_t spatial_shift = shifted ? sol_spatial_shift : 0; + const int64_t time_shift = shifted ? sol_temporal_shift : 0; for (int64_t t = 0; t < num_frames; ++t) { - 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; + const float time = config.sol ? ((t + time_shift) % sol_block_size + 0.5f) / sol_block_size + : (config.patch_size_t == 1 ? 0.5f / 7.f : (t + 0.5f) / num_frames); + float pt = (2.f * time - 1.f) * two_pi; for (int64_t h = 0; h < height; ++h) { - float ph = (2.f * ((h + 0.5f) / height) - 1.f) * two_pi; + const float y = config.sol ? ((h + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (h + 0.5f) / height; + float ph = (2.f * y - 1.f) * two_pi; for (int64_t w = 0; w < width; ++w) { - float pw = (2.f * ((w + 0.5f) / width) - 1.f) * two_pi; + const float x = config.sol ? ((w + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (w + 0.5f) / width; + float pw = (2.f * x - 1.f) * two_pi; ids.push_back({pt, ph, pw}); } } } - for (int i = 0; i < 5; ++i) { + for (int64_t i = 0; i < global_rows; ++i) { ids.push_back({0.f, 0.f, 0.f}); } auto values = Rope::embed_nd(ids, @@ -835,24 +988,49 @@ namespace MiniMaxH3VAE { const sd::Tensor& z, bool decode_graph) override { auto input = ensure_video_shape(z); + std::unique_ptr sol_plan; + std::unique_ptr shifted_plan; + sd::Tensor shifted_rope; + sd::Tensor sol_row_ids; if (decode_graph) { rope_cache = build_rope(input.shape()[0], input.shape()[1], input.shape()[2]); + if (config.sol) { + const auto& shape = input.shape(); + sol_plan = std::make_unique(shape[0], shape[1], shape[2], false); + shifted_plan = std::make_unique(shape[0], shape[1], shape[2], true); + shifted_rope = build_rope(shape[0], shape[1], shape[2], true); + const int64_t rows = std::max(shifted_rope.shape()[3], pixel_mean.shape()[3] * + sol_output_frames(shape[2], config.patch_size_t)); + sol_row_ids = sd::Tensor({rows}); + for (int64_t i = 0; i < rows; ++i) { + sol_row_ids.data()[i] = static_cast(i); + } + } } auto get_graph = [&]() -> ggml_cgraph* { auto value = make_input(input); auto mean = make_input(pixel_mean); auto std = make_input(pixel_std); - auto runner_ctx = get_context(); + auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE); + auto runner_ctx = get_context(graph); ggml_tensor* out = nullptr; if (decode_graph) { - auto pe = make_input(rope_cache); - out = model.decode(&runner_ctx, value, pe, mean, std); + auto pe = make_input(rope_cache); + auto shifted_pe = config.sol ? make_input(shifted_rope) : nullptr; + std::array sol_inputs; + if (sol_plan) { + auto bind_input = [&](const auto& tensor) { return make_input(tensor); }; + sol_inputs = {sol_plan->bind(bind_input), shifted_plan->bind(bind_input)}; + auto row_ids = make_input(sol_row_ids); + sol_inputs[0].row_ids = row_ids; + sol_inputs[1].row_ids = row_ids; + } + out = model.decode(&runner_ctx, value, pe, shifted_pe, sol_inputs, mean, std); } else { out = model.encode(&runner_ctx, value, mean, std); } - auto graph = new_graph_custom(H3_VIDEO_VAE_GRAPH_SIZE); ggml_build_forward_expand(graph, out); return graph; }; diff --git a/src/model/vae/vae_tiling.hpp b/src/model/vae/vae_tiling.hpp index acad046236..c88ff9f056 100644 --- a/src/model/vae/vae_tiling.hpp +++ b/src/model/vae/vae_tiling.hpp @@ -116,14 +116,14 @@ inline sd::Tensor process_vae_temporal_tiles(const sd::Tensor& inp return output; } -template +template inline sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor& input, const VAETemporalTilePlan& plan, - int output_scale, + FrameCount frame_count, + FrameOffset frame_offset, Fn&& on_processing) { - GGML_ASSERT(output_scale >= 1); - const int64_t output_frames = 1 + (input.shape()[2] - 1) * output_scale; - const int overlap_frames = plan.overlap > 0 ? 1 + (plan.overlap - 1) * output_scale : 0; + const int64_t output_frames = frame_count(input.shape()[2]); + const int64_t overlap_frames = frame_count(plan.overlap); std::vector weights(static_cast(output_frames), 0.f); sd::Tensor output; @@ -138,7 +138,7 @@ inline sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor +inline sd::Tensor process_vae_temporal_tiles_blended(const sd::Tensor& input, + const VAETemporalTilePlan& plan, + int output_scale, + Fn&& on_processing) { + GGML_ASSERT(output_scale >= 1); + auto frame_count = [&](int64_t frames) { return frames > 0 ? 1 + (frames - 1) * output_scale : 0; }; + auto frame_offset = [&](int64_t start) { return start * output_scale; }; + return process_vae_temporal_tiles_blended(input, plan, frame_count, frame_offset, std::forward(on_processing)); +} + #endif // __SD_MODEL_VAE_VAE_TILING_HPP__ From 56070ea0342e0bb93da55c858e952da4543110f9 Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 10 Nov 2026 23:38:47 +0800 Subject: [PATCH 52/57] refactor: replace QKV split helpers with ggml_ext_chunk (#2322) --- src/core/ggml_extend.cpp | 30 ------------------------------ src/core/ggml_extend.h | 8 -------- src/model/common/ggml_block.hpp | 2 +- src/model/diffusion/hunyuan.hpp | 2 +- src/model/diffusion/ideogram4.hpp | 2 +- src/model/diffusion/iris.hpp | 2 +- src/model/diffusion/lens.hpp | 4 ++-- src/model/diffusion/mmdit.hpp | 2 +- src/model/diffusion/pid.hpp | 6 +++--- src/model/te/llm.hpp | 2 +- src/model/vae/wan_vae.hpp | 2 +- 11 files changed, 12 insertions(+), 50 deletions(-) diff --git a/src/core/ggml_extend.cpp b/src/core/ggml_extend.cpp index 97561581f1..52d9cfc83f 100644 --- a/src/core/ggml_extend.cpp +++ b/src/core/ggml_extend.cpp @@ -541,36 +541,6 @@ ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx, return x; } -std::vector ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) { - auto ctx = runner_ctx->ggml_ctx; - qkv = ggml_reshape_4d(ctx, qkv, qkv->ne[0] / 3, 3, qkv->ne[1], qkv->ne[2]); // [N, L, 3, C] - qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 3, 1, 2)); // [3, N, L, C] - - int64_t offset = qkv->nb[2] * qkv->ne[2]; - auto q = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 0); // [N, L, C] - auto k = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 1); // [N, L, C] - auto v = ggml_view_3d(ctx, qkv, qkv->ne[0], qkv->ne[1], qkv->ne[2], qkv->nb[1], qkv->nb[2], offset * 2); // [N, L, C] - return {q, k, v}; -} - -std::vector ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv) { - auto ctx = runner_ctx->ggml_ctx; - int64_t W = qkv->ne[0]; - int64_t H = qkv->ne[1]; - int64_t C = qkv->ne[2] / 3; - int64_t N = qkv->ne[3]; - int64_t nb1 = qkv->nb[1]; - int64_t nb2 = qkv->nb[2]; - qkv = ggml_reshape_4d(ctx, qkv, W * H, C, 3, N); // [N, 3, C, H*W] - qkv = ggml_cont(ctx, ggml_ext_torch_permute(runner_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[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}; -} - ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx, float value, int64_t ne0, diff --git a/src/core/ggml_extend.h b/src/core/ggml_extend.h index f8c14c0ee6..1f0a42c292 100644 --- a/src/core/ggml_extend.h +++ b/src/core/ggml_extend.h @@ -150,14 +150,6 @@ ggml_tensor* ggml_ext_conv_3d(GGMLRunnerContext* runner_ctx, float scale = 1.f, bool allow_conv2d = false); -// qkv: [N, L, 3*C] -// return: ([N, L, C], [N, L, C], [N, L, C]) -std::vector ggml_ext_split_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv); - -// qkv: [N, 3*C, H, W] -// return: ([N, C, H, W], [N, C, H, W], [N, C, H, W]) -std::vector ggml_ext_split_image_qkv(GGMLRunnerContext* runner_ctx, ggml_tensor* qkv); - // Constant and cast helpers require the built-in tensors initialized by GGMLRunner. ggml_tensor* ggml_ext_full(GGMLRunnerContext* runner_ctx, float value, diff --git a/src/model/common/ggml_block.hpp b/src/model/common/ggml_block.hpp index eb437e5df5..805b8682d0 100644 --- a/src/model/common/ggml_block.hpp +++ b/src/model/common/ggml_block.hpp @@ -910,7 +910,7 @@ class MultiheadAttention : public GGMLBlock { if (proj_in) { auto in_proj = std::dynamic_pointer_cast(blocks[in_proj_name]); auto qkv = in_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); q = qkv_vec[0]; k = qkv_vec[1]; v = qkv_vec[2]; diff --git a/src/model/diffusion/hunyuan.hpp b/src/model/diffusion/hunyuan.hpp index d2bc8b5804..15c383de2a 100644 --- a/src/model/diffusion/hunyuan.hpp +++ b/src/model/diffusion/hunyuan.hpp @@ -50,7 +50,7 @@ namespace Hunyuan { // self attn auto qkv = self_attn_qkv->forward(ctx, norm1->forward(ctx, txt)); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); auto q = qkv_vec[0]; auto k = qkv_vec[1]; auto v = qkv_vec[2]; diff --git a/src/model/diffusion/ideogram4.hpp b/src/model/diffusion/ideogram4.hpp index fbc189aa3b..8d7cf2ae23 100644 --- a/src/model/diffusion/ideogram4.hpp +++ b/src/model/diffusion/ideogram4.hpp @@ -222,7 +222,7 @@ namespace Ideogram4 { auto out_proj = std::dynamic_pointer_cast(blocks["o"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, n_token, N); auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, n_token, N); auto v = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[2], head_dim, num_heads, n_token, N); diff --git a/src/model/diffusion/iris.hpp b/src/model/diffusion/iris.hpp index 92f4b2e207..662704e237 100644 --- a/src/model/diffusion/iris.hpp +++ b/src/model/diffusion/iris.hpp @@ -92,7 +92,7 @@ namespace Iris { ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe = nullptr, ggml_tensor* mask = nullptr) { auto c = ctx->ggml_ctx; - auto qkv = ggml_ext_split_qkv(ctx, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x)); + auto qkv = ggml_ext_chunk(ctx, std::dynamic_pointer_cast(blocks["qkv"])->forward(ctx, x), 3, 0); auto q = ggml_reshape_4d(c, qkv[0], dim / heads, heads, x->ne[1], x->ne[2]); auto k = ggml_reshape_4d(c, qkv[1], dim / heads, heads, x->ne[1], x->ne[2]); auto v = ggml_reshape_4d(c, qkv[2], dim / heads, heads, x->ne[1], x->ne[2]); diff --git a/src/model/diffusion/lens.hpp b/src/model/diffusion/lens.hpp index 5c86e65d01..92e81d31d3 100644 --- a/src/model/diffusion/lens.hpp +++ b/src/model/diffusion/lens.hpp @@ -151,8 +151,8 @@ namespace Lens { int64_t N = img->ne[2]; int64_t inner = dim_head * num_heads; - auto img_qkv_vec = ggml_ext_split_qkv(ctx, img_qkv->forward(ctx, img)); - auto txt_qkv_vec = ggml_ext_split_qkv(ctx, txt_qkv->forward(ctx, txt)); + auto img_qkv_vec = ggml_ext_chunk(ctx, img_qkv->forward(ctx, img), 3, 0); + auto txt_qkv_vec = ggml_ext_chunk(ctx, txt_qkv->forward(ctx, txt), 3, 0); auto img_q = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[0], dim_head, num_heads, n_img, N); auto img_k = ggml_reshape_4d(ctx->ggml_ctx, img_qkv_vec[1], dim_head, num_heads, n_img, N); diff --git a/src/model/diffusion/mmdit.hpp b/src/model/diffusion/mmdit.hpp index e409a01f0a..ea653273bf 100644 --- a/src/model/diffusion/mmdit.hpp +++ b/src/model/diffusion/mmdit.hpp @@ -334,7 +334,7 @@ class SelfAttention : public GGMLBlock { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); int64_t head_dim = qkv_vec[0]->ne[0] / num_heads; auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); // [N, n_token, n_head, d_head] auto k = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[1], head_dim, num_heads, qkv_vec[1]->ne[1], qkv_vec[1]->ne[2]); // [N, n_token, n_head, d_head] diff --git a/src/model/diffusion/pid.hpp b/src/model/diffusion/pid.hpp index 225a30aeb8..372e15729c 100644 --- a/src/model/diffusion/pid.hpp +++ b/src/model/diffusion/pid.hpp @@ -293,7 +293,7 @@ namespace Pid { auto proj = std::dynamic_pointer_cast(blocks["proj"]); auto qkv = qkv_proj->forward(ctx, x); - auto qkv_vec = ggml_ext_split_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); int64_t L = x->ne[1]; int64_t N = x->ne[2]; int64_t head_dim = dim / num_heads; @@ -343,14 +343,14 @@ namespace Pid { int64_t N = x->ne[2]; int64_t head_dim = dim / num_heads; - auto qkv_x = ggml_ext_split_qkv(ctx, qkv_x_proj->forward(ctx, x)); + auto qkv_x = ggml_ext_chunk(ctx, qkv_x_proj->forward(ctx, x), 3, 0); auto qx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[0], head_dim, num_heads, Nx, N); auto kx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[1], head_dim, num_heads, Nx, N); auto vx = ggml_reshape_4d(ctx->ggml_ctx, qkv_x[2], head_dim, num_heads, Nx, N); qx = q_norm_x->forward(ctx, qx); kx = k_norm_x->forward(ctx, kx); - auto qkv_y = ggml_ext_split_qkv(ctx, qkv_y_proj->forward(ctx, y)); + auto qkv_y = ggml_ext_chunk(ctx, qkv_y_proj->forward(ctx, y), 3, 0); auto qy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[0], head_dim, num_heads, Ny, N); auto ky = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[1], head_dim, num_heads, Ny, N); auto vy = ggml_reshape_4d(ctx->ggml_ctx, qkv_y[2], head_dim, num_heads, Ny, N); diff --git a/src/model/te/llm.hpp b/src/model/te/llm.hpp index 7ce1f175e3..e6d5626029 100644 --- a/src/model/te/llm.hpp +++ b/src/model/te/llm.hpp @@ -1252,7 +1252,7 @@ namespace LLM { } else { auto qkv_proj = std::dynamic_pointer_cast(blocks["qkv"]); auto qkv = qkv_proj->forward(ctx, x); - qkv_vec = ggml_ext_split_qkv(ctx, qkv); + qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 0); } auto q = ggml_reshape_4d(ctx->ggml_ctx, qkv_vec[0], head_dim, num_heads, qkv_vec[0]->ne[1], qkv_vec[0]->ne[2]); // [N, n_token, n_head, d_head] diff --git a/src/model/vae/wan_vae.hpp b/src/model/vae/wan_vae.hpp index 52b09bd786..d44fe92b40 100644 --- a/src/model/vae/wan_vae.hpp +++ b/src/model/vae/wan_vae.hpp @@ -623,7 +623,7 @@ namespace WAN { const int64_t w = x->ne[0]; auto qkv = to_qkv->forward(ctx, x); - auto qkv_vec = ggml_ext_split_image_qkv(ctx, qkv); + auto qkv_vec = ggml_ext_chunk(ctx, qkv, 3, 2); auto q = qkv_vec[0]; q = ggml_ext_cont(ctx, ggml_ext_torch_permute(ctx, q, 2, 0, 1, 3)); // [t, h, w, c] From ddc367ab5d502113e0df3b5d9415112ce89fac24 Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 10 Nov 2026 23:45:01 +0800 Subject: [PATCH 53/57] refactor: unify runner execution and weight residency (#2321) --- docs/backend.md | 24 +- src/core/execution_engine.cpp | 992 +++++++++++++++++++++++++ src/core/execution_engine.h | 107 +++ src/core/ggml_runner.cpp | 818 +------------------- src/core/ggml_runner.h | 76 +- src/device_residency_manager.h | 57 +- src/model/diffusion/qwen_image_2_1.hpp | 2 +- src/model_manager.cpp | 361 +++++---- src/model_manager.h | 30 +- src/model_manager_files.cpp | 42 +- src/pipeline/image.cpp | 22 +- src/pipeline/image_upscale.cpp | 138 ++++ src/pipeline/image_upscale.h | 57 ++ src/upscaler.cpp | 173 +---- src/upscaler.h | 9 +- 15 files changed, 1667 insertions(+), 1241 deletions(-) create mode 100644 src/core/execution_engine.cpp create mode 100644 src/core/execution_engine.h create mode 100644 src/pipeline/image_upscale.cpp create mode 100644 src/pipeline/image_upscale.h diff --git a/docs/backend.md b/docs/backend.md index da90e48968..cc6545da92 100644 --- a/docs/backend.md +++ b/docs/backend.md @@ -160,10 +160,12 @@ the runner's graph-cut capacity checks. Runtime capacity checks also leave 512 MiB of currently free device memory for backend scratch buffers and pipelines, including with explicit backend assignments. They cap free-memory reports by the device's total memory minus tracked -resident allocations. Vulkan reports exceeding total memory are rejected because -its heap-budget subtraction can underflow. Other backends use the cap instead of -treating such reports as zero free memory. Failed checks log the reported free and -total memory alongside tracked weight and runtime allocations. +resident allocations. Vulkan reports exceeding total memory are rejected on +discrete GPUs because its heap-budget subtraction can underflow. On integrated +GPUs, a wrapped report triggers reclamation before the tracked estimate can be +accepted, subject to the byte checks and configured budgets. Other backends use +the cap instead of treating such reports as zero free memory. Failed checks log +the reported free and total memory alongside tracked weight and runtime allocations. With `--mmap`, device-backed mappings count toward these budgets at their full mapped-file size, once per device buffer even when multiple parameter blocks share it. Mappings retained in the loader cache continue to count. @@ -181,8 +183,18 @@ GPU cache space follows the same component priority. Before a lower-priority component can become permanently resident, the planner leaves room for the full weights and estimated compute space of higher-priority offloaded components. If offloaded diffusion already needs the entire main GPU budget, TE and VAE also -use offloaded parameters. Their GPU copies can then be released after their -phases, leaving more room to reuse diffusion weights across sampling steps. +use offloaded parameters. GPU copies remain cached across sampling steps while +a component's execution is active and are released under memory pressure. +Ending a component's execution releases its staged GPU copies, disk-backed +weight buffers, temporary compute buffers and execution caches. Parameters +assigned permanently to a backend remain resident. Hires model upscaling shares +the runtime budget and can +reclaim idle weight copies from other components. +Reclamation releases idle compute buffers first, then cached weight allocations +until the current work fits. Weights used by active computation or pending +transfers stay protected. +For CPU computation with disk-backed parameters, cached weights are released +when execution ends so they do not accumulate in RAM across components. CPU parameter residency allows GPU weight caching; it does not force every weight to be copied again at every step. diff --git a/src/core/execution_engine.cpp b/src/core/execution_engine.cpp new file mode 100644 index 0000000000..5ad241a1cd --- /dev/null +++ b/src/core/execution_engine.cpp @@ -0,0 +1,992 @@ +#include "core/execution_engine.h" + +#include +#include +#include +#include +#include + +#include "core/ggml_extend_backend.h" +#include "core/ggml_tensor_utils.h" +#include "core/layer_split_partition.h" +#include "core/segment_graph_bindings.h" +#include "core/segment_weight_pipeline.h" +#include "core/tensor_ggml.hpp" +#include "core/util.h" + +std::shared_ptr DeviceResidencyManager::execution_engine() { + if (execution_engine_ == nullptr) { + execution_engine_ = std::make_shared(this); + } + return execution_engine_; +} + +namespace sd { + class GraphExecution { + ExecutionEngine& engine_; + ExecutionSession& session_; + const ComputeTask& task_; + std::shared_ptr manager_; + std::unordered_set params_tensor_set_; + ggml_status last_compute_status_ = GGML_STATUS_FAILED; + Tensor output_; + + using GraphCutPlan = ggml_graph_cut::Plan; + ggml_tensor* canonical_param_tensor(ggml_tensor* tensor); + std::vector collect_used_param_tensors(ggml_cgraph* graph); + ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const; + // Weightless ops follow the most recent weight device. Views must stay + // unpinned so the scheduler preserves cross-device copies for consumers. + void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* graph, ggml_cgraph* original_graph = nullptr); + Tensor read_graph_tensor(ggml_tensor* tensor, const char* label); + void copy_data_to_backend_tensor(ggml_cgraph* graph); + const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* graph); + bool assign_graph_cut_layer_split_backends(ggml_cgraph* graph); + 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, const std::function& on_before_compute); + bool execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs); + bool is_multi_device() const { return !session_.extra_backends_.empty(); } + size_t retained_runtime_buffer_bytes(ggml_backend_t backend) const { return session_.resident_bytes(backend); } + void sync_runtime_residency() { engine_.sync_residency(session_); } + + public: + GraphExecution(ExecutionEngine& engine, ExecutionSession& session, const ComputeTask& task) + : engine_(engine), session_(session), task_(task), manager_(session.residency_manager_.lock()) { + params_tensor_set_.insert(task.weights.begin(), task.weights.end()); + if (manager_ != nullptr) { + for (int i = 0; i < ggml_graph_cut::leaf_count(task.graph); ++i) { + auto* parameter = manager_->resolve_param_tensor(ggml_graph_cut::leaf_tensor(task.graph, i)); + if (parameter != nullptr) { + params_tensor_set_.insert(parameter); + } + } + for (auto* tensor : params_tensor_set_) { + auto* parameter = manager_->resolve_param_tensor(const_cast(tensor)); + if (parameter != nullptr) { + session_.params_.insert(parameter); + } + } + } + } + + ComputeResult run() { + if (!execute_graph(task_.graph, task_.n_threads, task_.no_return, task_.read_outputs)) { + return {last_compute_status_, {}}; + } + return {GGML_STATUS_SUCCESS, std::move(output_)}; + } + }; + + static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { + // Vulkan reports a preferred allocation chunk size, not a hard tensor limit. + if (sd_backend_is(backend, "Vulkan")) { + return true; + } + auto buft = ggml_backend_get_default_buffer_type(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; + if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { + return false; + } + // Cross-backend copies must also fit the destination's allocation limit. + for (auto src : node->src) { + if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { + return false; + } + } + return true; + } + + static size_t add_bytes(size_t a, size_t b) { + return b > SIZE_MAX - a ? SIZE_MAX : a + b; + } + + ggml_tensor* GraphExecution::canonical_param_tensor(ggml_tensor* tensor) { + for (auto* current = tensor; current != nullptr; current = current->view_src) { + if (params_tensor_set_.count(current) != 0) + return current; + } + return nullptr; + } + + std::vector GraphExecution::collect_used_param_tensors(ggml_cgraph* gf) { + std::vector used_params; + if (gf == nullptr || params_tensor_set_.empty()) { + return used_params; + } + + std::unordered_set seen_params; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + seen_params.reserve(static_cast(n_leafs)); + for (int i = 0; i < n_leafs; ++i) { + ggml_tensor* leaf = sd::ggml_graph_cut::leaf_tensor(gf, i); + ggml_tensor* param_leaf = canonical_param_tensor(leaf); + if (param_leaf != nullptr && + seen_params.insert(param_leaf).second) { + used_params.push_back(param_leaf); + } + } + return used_params; + } + + ggml_backend_t GraphExecution::backend_for_weight(const ggml_tensor* tensor) const { + if (tensor == nullptr || tensor->buffer == nullptr) { + return nullptr; + } + if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || + ggml_backend_buffer_is_host(tensor->buffer)) { + return nullptr; + } + ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); + if (dev == nullptr) { + return nullptr; + } + if (ggml_backend_get_device(session_.backend_) == dev) { + return session_.backend_; + } + for (ggml_backend_t backend : session_.extra_backends_) { + if (ggml_backend_get_device(backend) == dev) { + return backend; + } + } + return nullptr; + } + + void GraphExecution::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { + if (sched == nullptr || gf == nullptr) { + return; + } + ggml_backend_t current = session_.backend_; + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + auto node_assignment = session_.node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i)); + if (node_assignment != session_.node_assignments_.end()) { + current = node_assignment->second; + } + for (int s = 0; s < GGML_MAX_SRC; s++) { + ggml_backend_t weight_backend = backend_for_weight(node->src[s]); + if (weight_backend != nullptr) { + if (node_assignment == session_.node_assignments_.end()) { + current = weight_backend; + } + } + } + if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { + continue; + } + if (!backend_can_allocate_op(current, node)) { + auto cpu = session_.workspace_.cpu_backend(); + if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, cpu); + } + } else if (ggml_backend_supports_op(current, node)) { + ggml_backend_sched_set_tensor_backend(sched, node, current); + } + } + } + + Tensor GraphExecution::read_graph_tensor(ggml_tensor* tensor, const char* label) { + if (tensor == nullptr) { + LOG_ERROR("%s %s tensor is null", task_.description.c_str(), label); + return {}; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_ERROR("%s %s tensor type mismatch: got %s", + task_.description.c_str(), + label, + ggml_type_name(tensor->type)); + return {}; + } + ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); + if (buf == nullptr) { + LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", + task_.description.c_str(), + label, + tensor->name[0] != '\0' ? tensor->name : "", + ggml_op_name(tensor->op), + tensor->buffer, + tensor->view_src, + tensor->view_src ? tensor->view_src->buffer : nullptr, + tensor->data); + return {}; + } + + return sd::make_sd_tensor_from_ggml(tensor); + } + + void GraphExecution::copy_data_to_backend_tensor(ggml_cgraph* gf) { + GGML_ASSERT(gf != nullptr); + std::unordered_set graph_tensor_set; + const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); + const int n_nodes = ggml_graph_n_nodes(gf); + graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); + for (int i = 0; i < n_leafs; ++i) { + graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); + } + for (int i = 0; i < n_nodes; ++i) { + graph_tensor_set.insert(ggml_graph_node(gf, i)); + } + + for (auto& kv : task_.inputs) { + auto tensor = kv.first; + auto data = kv.second; + if (tensor == nullptr || data == nullptr) { + continue; + } + const char* name = ggml_get_name(tensor); + if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { + continue; + } + if (tensor->buffer == nullptr) { + LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", + task_.description.c_str(), + name != nullptr ? name : "", + (long long)tensor->ne[0], + (long long)tensor->ne[1], + (long long)tensor->ne[2], + (long long)tensor->ne[3], + ggml_type_name(tensor->type)); + continue; + } + + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + if (buf == nullptr) { + LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", + task_.description.c_str(), + tensor && tensor->name[0] != '\0' ? tensor->name : "", + tensor ? ggml_op_name(tensor->op) : "", + data, + tensor ? tensor->view_src : nullptr, + (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); + continue; + } + + ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); + } + } + + const GraphExecution::GraphCutPlan& GraphExecution::resolve_graph_cut_plan(ggml_cgraph* gf) { + GGML_ASSERT(gf != nullptr); + return sd::ggml_graph_cut::resolve_plan(session_.backend_, + gf, + &session_.plan_cache_, + params_tensor_set_, + task_.description.c_str()); + } + + bool GraphExecution::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { + session_.node_assignments_.clear(); + if (!session_.layer_split_enabled_) { + return true; + } + if (!is_multi_device()) { + LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", task_.description.c_str()); + return false; + } + + const auto& plan = resolve_graph_cut_plan(gf); + if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { + auto manager = manager_; + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str()); + return false; + } + std::vector graph_params = collect_used_param_tensors(gf); + if (!graph_params.empty() && + !manager->assign_compute_backend(graph_params, session_.backend_)) { + LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(session_.backend_).c_str()); + return false; + } + for (ggml_tensor* param : graph_params) { + if (param != nullptr) { + session_.weight_assignments_[param] = session_.backend_; + } + } + const int n_nodes = ggml_graph_n_nodes(gf); + for (int i = 0; i < n_nodes; i++) { + ggml_tensor* node = ggml_graph_node(gf, i); + if (node != nullptr) { + session_.node_assignments_[node] = session_.backend_; + } + } + if (!session_.primary_notice_logged_) { + LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(session_.backend_).c_str(), + graph_params.size()); + session_.primary_notice_logged_ = true; + } else { + LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(session_.backend_).c_str(), + graph_params.size()); + } + return true; + } + + std::vector split_backends; + split_backends.reserve(session_.extra_backends_.size() + 1); + split_backends.push_back(session_.backend_); + for (ggml_backend_t backend : session_.extra_backends_) { + if (backend != nullptr) { + split_backends.push_back(backend); + } + } + + auto manager = manager_; + if (manager == nullptr) { + LOG_ERROR("%s weight manager is not set for graph-cut layer split", task_.description.c_str()); + return false; + } + + sd::GraphCutLayerSplitAssignment assignment; + auto canonicalize_param = [this](ggml_tensor* tensor) { + return canonical_param_tensor(tensor); + }; + if (!sd::partition_graph_cut_layer_split(task_.description.c_str(), + gf, + plan, + split_backends, + session_.backend_limits_, + session_.max_backend_bytes_, + session_.weight_assignments_, + canonicalize_param, + &assignment)) { + return false; + } + + for (size_t i = 0; i < split_backends.size(); i++) { + if (assignment.tensors_by_backend[i].empty()) { + continue; + } + if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { + LOG_ERROR("%s graph-cut layer split failed to assign params to %s", + task_.description.c_str(), + sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); + return false; + } + } + + session_.node_assignments_ = std::move(assignment.node_assignments); + sd::log_graph_cut_layer_split_assignment(task_.description.c_str(), split_backends, assignment); + + return true; + } + + ComputeWorkspace::Measurement GraphExecution::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 session_.backend_; + } + if (!params_tensor_set_.count(tensor)) { + return nullptr; + } + auto placement = session_.weight_assignments_.find(tensor); + return placement == session_.weight_assignments_.end() ? session_.backend_ : placement->second; + }; + auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { + pin_multi_device_nodes(scheduler, copy, graph); + }; + return session_.workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device()); + } + + std::vector GraphExecution::memory_requests( + const std::vector& sizes, + size_t pending_cache_bytes) const { + std::vector requests; + for (const auto& size : sizes) { + const size_t retained = retained_runtime_buffer_bytes(size.backend); + const size_t reusable = session_.workspace_.bytes(size.backend); + const size_t cache_bytes = size.backend == session_.backend_ ? pending_cache_bytes : 0; + const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes); + size_t limit = session_.max_backend_bytes_; + if (is_multi_device()) { + size_t index = 0; + if (size.backend != session_.backend_) { + auto position = std::find(session_.extra_backends_.begin(), session_.extra_backends_.end(), size.backend); + index = static_cast(position - session_.extra_backends_.begin()) + 1; + } + if (index < session_.backend_limits_.size()) { + limit = session_.backend_limits_[index]; + } + } + requests.push_back({size.backend, session_.owner_id(), pending, + retained, limit}); + } + return requests; + } + + bool GraphExecution::fits(const std::vector& requests, + const std::vector& params) const { + auto manager = manager_; + if (manager == nullptr) { + return params.empty(); + } + for (const auto& request : requests) { + if (!manager->fits_compute_backend_capacity(request, params)) { + return false; + } + } + return true; + } + + bool GraphExecution::execute_segment(ggml_cgraph* graph, int n_threads, const std::function& on_before_compute) { + if (sd_backend_is_cpu(session_.backend_)) { + sd_backend_cpu_set_n_threads(session_.backend_, n_threads); + } + if (session_.workspace_.cpu_backend() != nullptr) { + sd_backend_cpu_set_n_threads(session_.workspace_.cpu_backend(), n_threads); + } + auto scheduler = session_.workspace_.scheduler(); + 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{session_.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 + session_.workspace_.synchronize(); + if (status != GGML_STATUS_SUCCESS) { + last_compute_status_ = status; + LOG_ERROR("%s compute failed: %s", task_.description.c_str(), ggml_status_to_string(status)); + return false; + } + const std::string description = task_.description; + if (!task_.debug_tensors.empty()) { + std::unordered_set graph_tensors; + const int leaf_count = ggml_graph_cut::leaf_count(graph); + const int node_count = ggml_graph_n_nodes(graph); + graph_tensors.reserve(static_cast(leaf_count + node_count)); + for (int index = 0; index < leaf_count; ++index) { + graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index)); + } + for (int index = 0; index < node_count; ++index) { + graph_tensors.insert(ggml_graph_node(graph, index)); + } + + for (const auto& entry : task_.debug_tensors) { + ggml_tensor* tensor = entry.first; + if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) { + continue; + } + ggml_backend_buffer_t buffer = + tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; + if (buffer == nullptr) { + LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", + description.c_str(), + entry.second.c_str()); + continue; + } + if (tensor->type != GGML_TYPE_F32) { + LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", + description.c_str(), + entry.second.c_str(), + ggml_type_name(tensor->type)); + continue; + } + auto debug_tensor = make_sd_tensor_from_ggml(tensor); + print_sd_tensor(debug_tensor, false, entry.second.c_str()); + } + } + + return true; + } + + bool GraphExecution::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs) { + if (!assign_graph_cut_layer_split_backends(graph)) { + return false; + } + const auto params = collect_used_param_tensors(graph); + const auto& cached_plan = resolve_graph_cut_plan(graph); + auto manager = manager_; + const bool can_segment = !is_multi_device() && !sd_backend_is_cpu(session_.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({{session_.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 false; + } + 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, session_.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(); + auto& segment = monolithic_plan.segments.back(); + segment.group_name = "graph"; + segment.compute_buffer_size = cached_plan.compute_buffer_size; + segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph)); + segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph)); + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + segment.internal_node_indices.push_back(i); + } + for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { + auto tensor = ggml_graph_cut::leaf_tensor(graph, i); + ggml_graph_cut::Segment::InputRef input; + input.leaf_index = i; + input.type = canonical_param_tensor(tensor) != nullptr + ? ggml_graph_cut::Segment::INPUT_PARAM + : ggml_graph_cut::Segment::INPUT_EXTERNAL; + segment.input_refs.push_back(input); + } + } + const auto& plan = segmented ? cached_plan : monolithic_plan; + const bool segments_changed = plan.segments.size() != session_.logged_segment_count_; + if (segments_changed && (segmented || session_.logged_segment_count_ > 1)) { + LOG_VERBOSE("%s using %zu segment%s", task_.description.c_str(), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + SegmentGraphBindings bindings(session_.cut_cache_, plan, graph); + SegmentWeightPipeline weights(manager, session_.backend_, session_.owner_id(), + graph, plan, params_tensor_set_, + segmented && manager != nullptr && manager->prefetch_enabled()); + + std::map peak_compute_bytes; + auto track_compute_buffer = [&](ggml_backend_t backend) { + if (backend != nullptr) { + auto& peak = peak_compute_bytes[backend]; + peak = std::max(peak, session_.workspace_.bytes(backend)); + } + }; + for (size_t index = 0; index < plan.segments.size(); ++index) { + const auto& segment = plan.segments[index]; + const bool last = index + 1 == plan.segments.size(); + auto fail_segment = [&](const char* phase) { + LOG_ERROR("%s segment %zu/%zu (%s) failed during %s", task_.description.c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str(), phase); + return false; + }; + session_.cut_cache_.prune(segment.live_cut_names); + bindings.reset(segment); + if (!bindings.bind_cached_inputs(segment, task_.description.c_str())) { + return fail_segment("input binding"); + } + ggml_context* segment_context = nullptr; + auto segment_graph = segmented + ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context) + : graph; + struct SegmentCleanup { + GraphExecution& runner; + SegmentWeightPipeline& weights; + SegmentGraphBindings& bindings; + ggml_context* context; + ~SegmentCleanup() { + if (!runner.session_.workspace_.segment_end()) { + runner.last_compute_status_ = GGML_STATUS_FAILED; + } + bindings.restore(); + weights.segment_end(); + ggml_free(context); + runner.sync_runtime_residency(); + } + } segment_cleanup{*this, weights, bindings, segment_context}; + + 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(session_.backend_, tensor)) { + cache_outputs.insert(tensor); + } + } + } + auto measurement = segmented ? measure(segment_graph, &cache_outputs) : full_measurement; + if (!session_.workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + const size_t cut_bytes = last ? 0 : session_.cut_cache_.estimate_output_bytes(graph, segment); + size_t new_cache_bytes = add_bytes(cut_bytes, session_.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))) { + return true; + } + if (session_.workspace_.release_excess(measurement)) { + sync_runtime_residency(); + requests = memory_requests(measurement.buffers, new_cache_bytes); + } + const bool ready = weights.ensure_segment_capacity(index, requests); + if (!ready && manager != nullptr) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + } + return ready; + }; + if (!weights.segment_start(index, ensure_capacity)) { + return fail_segment("weight preparation"); + } + // Preparing weights can execute LoRA graphs and reclaim an idle workspace. + if (!session_.workspace_.prepare(measurement)) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace preparation"); + } + if (!ensure_capacity()) { + return fail_segment("workspace capacity check"); + } + if (!session_.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 : session_.cut_cache_.estimate_output_bytes(graph, segment), session_.cache_.pending_bytes(segment_graph)); + if (!session_.workspace_.allocate(segment_graph, [&](ggml_backend_sched_t scheduler, ggml_cgraph* current) { + pin_multi_device_nodes(scheduler, current); + })) { + last_compute_status_ = GGML_STATUS_ALLOC_FAILED; + return fail_segment("workspace allocation"); + } + for (const auto& size : measurement.buffers) { + track_compute_buffer(size.backend); + } + if (session_.workspace_.scheduler() != nullptr) { + track_compute_buffer(session_.workspace_.cpu_backend()); + } + if (!ensure_capacity()) { + return fail_segment("allocated capacity check"); + } + copy_data_to_backend_tensor(segment_graph); + auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes); + if (!prefetch_requests.empty()) { + weights.enqueue_next(index, prefetch_requests.front()); + } + LOG_DEBUG("%s executing segment %zu/%zu: %s", task_.description.c_str(), + index + 1, plan.segments.size(), segment.group_name.c_str()); + if (!execute_segment(segment_graph, n_threads, [&]() { weights.submit_next(); })) { + return fail_segment("execution"); + } + auto cache_status = session_.cache_.capture(segment_graph); + if (cache_status == GGML_STATUS_SUCCESS) { + cache_status = session_.cut_cache_.capture(graph, segment, task_.description.c_str()); + } + if (cache_status != GGML_STATUS_SUCCESS) { + last_compute_status_ = cache_status; + return fail_segment("output caching"); + } + sync_runtime_residency(); + if (last) { + if (read_outputs && !read_outputs()) { + return fail_segment("output finalization"); + } + if (!no_return) { + output_ = read_graph_tensor(task_.output, "output"); + if (output_.empty()) { + return fail_segment("output readback"); + } + } + } + if (!session_.workspace_.segment_end()) { + last_compute_status_ = GGML_STATUS_FAILED; + return fail_segment("workspace synchronization"); + } + // Final outputs and their callbacks may still be views of consumed cuts. + session_.cut_cache_.prune(segment.future_cut_names); + } + if (segments_changed || peak_compute_bytes != session_.logged_compute_bytes_) { + for (const auto& entry : peak_compute_bytes) { + LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)", + task_.description.c_str(), entry.second / (1024.0 * 1024.0), + sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first), + plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); + } + session_.logged_compute_bytes_ = std::move(peak_compute_bytes); + session_.logged_segment_count_ = plan.segments.size(); + } + return true; + } + + ExecutionSession::ExecutionSession(ggml_backend_t backend, const std::shared_ptr& manager) + : backend_(backend), residency_manager_(manager), cache_(backend), cut_cache_(backend), workspace_(backend) { + GGML_ASSERT(backend_ != nullptr); + } + + size_t ExecutionSession::resident_bytes(ggml_backend_t backend) const { + size_t bytes = workspace_.bytes(backend); + if (backend == backend_) { + const auto add = [](size_t a, size_t b) { return b > SIZE_MAX - a ? SIZE_MAX : a + b; }; + bytes = add(bytes, cache_.resident_bytes(ggml_backend_get_device(backend))); + bytes = add(bytes, cut_cache_.resident_bytes(ggml_backend_get_device(backend))); + } + return bytes; + } + + void ExecutionSession::set_layer_split_enabled(bool enabled) { + layer_split_enabled_ = enabled; + if (!enabled) { + weight_assignments_.clear(); + node_assignments_.clear(); + primary_notice_logged_ = false; + } + } + + void ExecutionSession::set_backend_limits(const std::vector& limits) { + backend_limits_ = limits; + weight_assignments_.clear(); + node_assignments_.clear(); + primary_notice_logged_ = false; + } + + void ExecutionSession::set_runtime_backends(const std::vector& backends) { + extra_backends_.clear(); + for (auto backend : backends) { + if (backend != nullptr && backend != backend_ && + std::find(extra_backends_.begin(), extra_backends_.end(), backend) == extra_backends_.end()) { + extra_backends_.push_back(backend); + } + } + workspace_.set_extra_backends(extra_backends_); + weight_assignments_.clear(); + node_assignments_.clear(); + primary_notice_logged_ = false; + } + + ExecutionEngine::~ExecutionEngine() { + for (auto& entry : sessions_) { + end(*entry.second); + } + } + + ExecutionSession& ExecutionEngine::create_session(ggml_backend_t backend, const std::shared_ptr& manager) { + GGML_ASSERT(manager.get() == manager_); + auto session = std::make_unique(backend, manager); + auto& result = *session; + sessions_.emplace(session->owner_id(), std::move(session)); + return result; + } + + void ExecutionEngine::release_session(ExecutionSession& session) { + end(session); + sessions_.erase(session.owner_id()); + } + + bool ExecutionEngine::start(ExecutionSession& session) { + if (session.started_) { + return true; + } + sd_backend_set_buffer_cache_enabled(session.backend_, true); + for (auto backend : session.extra_backends_) { + sd_backend_set_buffer_cache_enabled(backend, true); + } + session.workspace_.set_extra_backends(session.extra_backends_); + if (auto manager = session.residency_manager_.lock()) { + manager->register_runtime_owner(session.owner_id()); + } + session.started_ = true; + return true; + } + + void ExecutionEngine::end(ExecutionSession& session) { + GGML_ASSERT(!session.executing_); + if (!session.started_) { + return; + } + session.workspace_.release(); + session.cache_.clear(); + session.cut_cache_.clear(); + session.logged_compute_bytes_.clear(); + session.logged_segment_count_ = 0; + if (auto manager = session.residency_manager_.lock()) { + manager->clear_prefetched_params(session.owner_id()); + std::vector params(session.params_.begin(), session.params_.end()); + manager->evict_compute_backend_params(params); + manager->remove_runtime_owner(session.owner_id()); + } + session.params_.clear(); + sd_backend_clear_buffer_cache(session.backend_); + for (auto backend : session.extra_backends_) { + sd_backend_clear_buffer_cache(backend); + } + session.started_ = false; + } + + void ExecutionEngine::sync_residency(ExecutionSession& session) { + if (auto manager = session.residency_manager_.lock()) { + manager->update_runtime_residency(session.owner_id(), session.backend_, session.resident_bytes(session.backend_)); + for (auto backend : session.extra_backends_) { + manager->update_runtime_residency(session.owner_id(), backend, session.resident_bytes(backend)); + } + } + } + + bool ExecutionEngine::ensure_capacity( + DeviceResidencyManager& manager, + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) { + GGML_ASSERT(&manager == manager_); + auto capacity = [&]() { return manager.inspect_compute_backend_capacity(request, required_params); }; + const auto initial = capacity(); + if (!initial.valid) { + return false; + } + if (initial.fits()) { + return true; + } + for (auto& entry : sessions_) { + auto& session = *entry.second; + if (session.owner_id() != request.owner_id && session.workspace_.release()) { + sync_residency(session); + if (capacity().fits(true)) { + return true; + } + } + } + auto protected_weights = protected_params; + protected_weights.insert(protected_weights.end(), required_params.begin(), required_params.end()); + auto candidates = manager.evictable_compute_backend_blocks(request.compute_backend, protected_weights); + std::unordered_map preferred_positions; + for (size_t index = 0; index < preferred_eviction_order.size(); ++index) { + for (auto* tensor : preferred_eviction_order[index]) { + auto* weight = manager.resolve_param_tensor(tensor); + preferred_positions.emplace(weight != nullptr ? weight : tensor, index); + } + } + auto preferred_position = [&](const WeightResidencyBlock& block) { + size_t position = 0; + // A shared allocation follows the last of its members in the eviction order. + for (auto* tensor : block.tensors) { + const auto found = preferred_positions.find(tensor); + if (found == preferred_positions.end()) { + return SIZE_MAX; + } + position = std::max(position, found->second); + } + return position; + }; + std::stable_sort(candidates.begin(), candidates.end(), [&](const auto& lhs, const auto& rhs) { + if (lhs.requires_demotion != rhs.requires_demotion) { + return !lhs.requires_demotion; + } + if (!lhs.requires_demotion) { + const auto left = preferred_position(lhs), right = preferred_position(rhs); + if (left != right) { + return left < right; + } + if (left != SIZE_MAX) { + return false; + } + } + return lhs.last_use_epoch < rhs.last_use_epoch; + }); + for (const auto& block : candidates) { + if (manager.evict_compute_backend_block(block) && capacity().fits(true)) { + return true; + } + } + auto available = capacity(); + if (available.fits(true)) { + return true; + } + if (sd_backend_is(request.compute_backend, "Vulkan") && + available.required_budget_bytes <= available.available_budget_bytes) { + ggml_backend_synchronize(request.compute_backend); + sd_backend_clear_buffer_cache(request.compute_backend); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(100); + do { + available = capacity(); + if (available.fits(true)) { + return true; + } + if (available.required_budget_bytes > available.available_budget_bytes) { + break; + } + std::this_thread::sleep_for(std::chrono::milliseconds(10)); + } while (std::chrono::steady_clock::now() < deadline); + } + available = manager.inspect_compute_backend_capacity(request, required_params, true); + const auto device_bytes = available.available_device_bytes == SIZE_MAX + ? std::string("unknown") + : sd_format("%.2f MB", available.available_device_bytes / (1024.0 * 1024.0)); + const auto budget_bytes = available.available_budget_bytes == SIZE_MAX + ? std::string("unlimited") + : sd_format("%.2f MB", available.available_budget_bytes / (1024.0 * 1024.0)); + LOG_WARN("execution engine cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget", + ggml_backend_name(request.compute_backend), + available.required_device_bytes / (1024.0 * 1024.0), + available.required_budget_bytes / (1024.0 * 1024.0), device_bytes.c_str(), budget_bytes.c_str()); + return false; + } + + void ExecutionEngine::clear_cache(ExecutionSession& session) { + session.cache_.clear(); + sync_residency(session); + } + + void ExecutionEngine::discard_graph(ExecutionSession& session) { + GGML_ASSERT(!session.executing_); + session.cache_.graph_end(false); + session.cut_cache_.clear(); + session.workspace_.release(); + sync_residency(session); + } + + ComputeResult ExecutionEngine::execute(ExecutionSession& session, const ComputeTask& task) { + if (session.executing_ || task.graph == nullptr || !start(session)) { + return {}; + } + session.executing_ = true; + ComputeResult result; + try { + GraphExecution execution(*this, session, task); + result = execution.run(); + } catch (const std::bad_alloc&) { + result.status = GGML_STATUS_ALLOC_FAILED; + LOG_ERROR("%s graph allocation failed", task.description.c_str()); + } catch (const std::exception& error) { + LOG_ERROR("%s graph execution failed on %s: %s", task.description.c_str(), + ggml_backend_name(session.backend_), error.what()); + } catch (...) { + LOG_ERROR("%s graph execution failed on %s: unknown exception", task.description.c_str(), + ggml_backend_name(session.backend_)); + } + if (!session.workspace_.segment_end()) { + result.status = GGML_STATUS_FAILED; + } + const bool success = result.status == GGML_STATUS_SUCCESS; + session.cache_.graph_end(success); + session.cut_cache_.clear(); + if (!success) { + result.output = {}; + session.workspace_.release(); + } + session.executing_ = false; + sync_residency(session); + return result; + } +} diff --git a/src/core/execution_engine.h b/src/core/execution_engine.h new file mode 100644 index 0000000000..79ced1d087 --- /dev/null +++ b/src/core/execution_engine.h @@ -0,0 +1,107 @@ +#ifndef __SD_CORE_EXECUTION_ENGINE_H__ +#define __SD_CORE_EXECUTION_ENGINE_H__ + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "core/compute_workspace.h" +#include "core/ggml_graph_cut.h" +#include "core/runner_cache.h" +#include "core/tensor.hpp" +#include "device_residency_manager.h" + +namespace sd { + struct ComputeTask { + std::string description; + ggml_cgraph* graph = nullptr; + ggml_tensor* output = nullptr; + std::vector weights; + std::map inputs; + std::vector> debug_tensors; + int n_threads = 1; + bool no_return = false; + std::function read_outputs; + }; + + struct ComputeResult { + ggml_status status = GGML_STATUS_FAILED; + Tensor output; + }; + + class ExecutionSession { + friend class ExecutionEngine; + friend class GraphExecution; + + ggml_backend_t backend_; + std::weak_ptr residency_manager_; + RunnerCache cache_; + GraphCutTensorCache cut_cache_; + ComputeWorkspace workspace_; + std::unordered_set params_; + std::vector extra_backends_; + size_t max_backend_bytes_ = 0; + bool layer_split_enabled_ = false; + std::vector backend_limits_; + ggml_graph_cut::PlanCache plan_cache_; + std::unordered_map weight_assignments_; + std::unordered_map node_assignments_; + bool primary_notice_logged_ = false; + std::map logged_compute_bytes_; + size_t logged_segment_count_ = 0; + bool started_ = false; + bool executing_ = false; + + size_t resident_bytes(ggml_backend_t backend) const; + uintptr_t owner_id() const { return reinterpret_cast(this); } + + public: + ExecutionSession(ggml_backend_t backend, const std::shared_ptr& manager); + ExecutionSession(const ExecutionSession&) = delete; + ExecutionSession& operator=(const ExecutionSession&) = delete; + + RunnerCache& cache() { return cache_; } + bool started() const { return started_; } + size_t compute_bytes() const { return workspace_.bytes(backend_); } + void clear_plan_cache() { plan_cache_.graph_cut_plans.clear(); } + void set_max_backend_bytes(size_t bytes) { max_backend_bytes_ = bytes; } + void set_layer_split_enabled(bool enabled); + void set_backend_limits(const std::vector& limits); + void set_runtime_backends(const std::vector& backends); + }; + + class ExecutionEngine { + friend class GraphExecution; + + std::map> sessions_; + DeviceResidencyManager* manager_; + + void sync_residency(ExecutionSession& session); + + public: + explicit ExecutionEngine(DeviceResidencyManager* manager = nullptr) + : manager_(manager) {} + ~ExecutionEngine(); + ExecutionSession& create_session(ggml_backend_t backend, const std::shared_ptr& manager); + void release_session(ExecutionSession& session); + bool start(ExecutionSession& session); + void end(ExecutionSession& session); + void clear_cache(ExecutionSession& session); + void discard_graph(ExecutionSession& session); + bool ensure_capacity(DeviceResidencyManager& manager, + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params); + // Graphs, input data and output callbacks remain borrowed until execute returns. + ComputeResult execute(ExecutionSession& session, const ComputeTask& task); + }; +} + +#endif // __SD_CORE_EXECUTION_ENGINE_H__ diff --git a/src/core/ggml_runner.cpp b/src/core/ggml_runner.cpp index 39d5429744..7a27811911 100644 --- a/src/core/ggml_runner.cpp +++ b/src/core/ggml_runner.cpp @@ -1,15 +1,11 @@ #include #include -#include #include #include "core/ggml_extend.h" #include "core/ggml_extend_backend.h" #include "core/ggml_runner.h" #include "core/ggml_tensor_utils.h" -#include "core/layer_split_partition.h" -#include "core/segment_graph_bindings.h" -#include "core/segment_weight_pipeline.h" using namespace sd; @@ -101,16 +97,6 @@ std::vector GGMLRunner::collect_used_param_tensors(ggml_cgraph* gf return used_params; } -void GGMLRunner::evict_compute_backend_param_tensors(const std::vector& tensors) { - if (tensors.empty()) { - return; - } - auto manager = residency_manager.lock(); - if (manager != nullptr) { - manager->evict_compute_backend_params(tensors); - } -} - void GGMLRunner::prepare_build_in_tensor_before() { one_tensor = ggml_new_tensor_1d(compute_ctx, GGML_TYPE_F32, 1); ggml_set_name(one_tensor, "ggml_runner_build_in_tensor:one"); @@ -148,7 +134,7 @@ ggml_cgraph* GGMLRunner::get_compute_graph(get_graph_cb_t get_graph) { ggml_build_forward_expand(gf, entry.first); } } - for (const auto& entry : cache_.outputs()) { + for (const auto& entry : execution_session_.cache().outputs()) { if (entry.second != nullptr) { ggml_build_forward_expand(gf, entry.second); } @@ -172,365 +158,22 @@ bool GGMLRunner::prepare_compute_graph(get_graph_cb_t get_graph, return true; } -ggml_backend_t GGMLRunner::backend_for_weight(const ggml_tensor* tensor) const { - if (tensor == nullptr || tensor->buffer == nullptr) { - return nullptr; - } - if (ggml_backend_buffer_get_usage(tensor->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS || - ggml_backend_buffer_is_host(tensor->buffer)) { - return nullptr; - } - ggml_backend_dev_t dev = ggml_backend_buft_get_device(ggml_backend_buffer_get_type(tensor->buffer)); - if (dev == nullptr) { - return nullptr; - } - if (ggml_backend_get_device(runtime_backend) == dev) { - return runtime_backend; - } - for (ggml_backend_t backend : extra_runtime_backends) { - if (ggml_backend_get_device(backend) == dev) { - return backend; - } - } - return nullptr; -} - -static bool backend_can_allocate_op(ggml_backend_t backend, const ggml_tensor* node) { - // Vulkan reports a preferred allocation chunk size, not a hard tensor limit. - if (sd_backend_is(backend, "Vulkan")) { - return true; - } - auto buft = ggml_backend_get_default_buffer_type(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; - if (ggml_backend_buft_get_alloc_size(buft, node) > max_size) { - return false; - } - // Cross-backend copies must also fit the destination's allocation limit. - for (auto src : node->src) { - if (src != nullptr && ggml_backend_buft_get_alloc_size(buft, src) > max_size) { - return false; - } - } - return true; -} - -void GGMLRunner::pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph) { - if (sched == nullptr || gf == nullptr) { - return; - } - ggml_backend_t current = runtime_backend; - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor* node = ggml_graph_node(gf, i); - auto node_assignment = graph_cut_layer_split_node_assignments_.find(original_graph == nullptr ? node : ggml_graph_node(original_graph, i)); - if (node_assignment != graph_cut_layer_split_node_assignments_.end()) { - current = node_assignment->second; - } - for (int s = 0; s < GGML_MAX_SRC; s++) { - ggml_backend_t weight_backend = backend_for_weight(node->src[s]); - if (weight_backend != nullptr) { - if (node_assignment == graph_cut_layer_split_node_assignments_.end()) { - current = weight_backend; - } - } - } - if (node->op == GGML_OP_NONE || node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE || - node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE) { - continue; - } - if (!backend_can_allocate_op(current, node)) { - auto cpu = workspace_.cpu_backend(); - if (cpu != nullptr && ggml_backend_supports_op(cpu, node)) { - ggml_backend_sched_set_tensor_backend(sched, node, cpu); - } - } else if (ggml_backend_supports_op(current, node)) { - ggml_backend_sched_set_tensor_backend(sched, node, current); - } - } -} - -size_t GGMLRunner::retained_runtime_buffer_bytes(ggml_backend_t backend) const { - backend = backend == nullptr ? runtime_backend : backend; - size_t bytes = workspace_.bytes(backend); - if (backend == runtime_backend) { - const size_t cache_bytes = cache_.resident_bytes(ggml_backend_get_device(backend)); - bytes = cache_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cache_bytes; - const size_t cut_bytes = cut_cache_.resident_bytes(ggml_backend_get_device(backend)); - bytes = cut_bytes > SIZE_MAX - bytes ? SIZE_MAX : bytes + cut_bytes; - } - return bytes; -} - -void GGMLRunner::sync_runtime_residency() { - if (auto manager = residency_manager.lock()) { - manager->update_runtime_residency(reinterpret_cast(this), - runtime_backend, retained_runtime_buffer_bytes()); - for (auto backend : extra_runtime_backends) { - manager->update_runtime_residency(reinterpret_cast(this), - backend, retained_runtime_buffer_bytes(backend)); - } - } -} - -std::optional> GGMLRunner::read_graph_tensor(ggml_tensor* tensor, const char* label) { - if (tensor == nullptr) { - LOG_ERROR("%s %s tensor is null", get_desc().c_str(), label); - return std::nullopt; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_ERROR("%s %s tensor type mismatch: got %s", - get_desc().c_str(), - label, - ggml_type_name(tensor->type)); - return std::nullopt; - } - ggml_backend_buffer_t buf = sd::ggml_graph_cut::tensor_buffer(tensor); - if (buf == nullptr) { - LOG_ERROR("%s %s tensor buffer missing: name=%s op=%s buffer=%p view_src=%p view_src_buffer=%p data=%p", - get_desc().c_str(), - label, - tensor->name[0] != '\0' ? tensor->name : "", - ggml_op_name(tensor->op), - tensor->buffer, - tensor->view_src, - tensor->view_src ? tensor->view_src->buffer : nullptr, - tensor->data); - return std::nullopt; - } - - return sd::make_sd_tensor_from_ggml(tensor); -} - -void GGMLRunner::copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy) { - GGML_ASSERT(gf != nullptr); - std::unordered_set graph_tensor_set; - const int n_leafs = sd::ggml_graph_cut::leaf_count(gf); - const int n_nodes = ggml_graph_n_nodes(gf); - graph_tensor_set.reserve(static_cast(n_leafs + n_nodes)); - for (int i = 0; i < n_leafs; ++i) { - graph_tensor_set.insert(sd::ggml_graph_cut::leaf_tensor(gf, i)); - } - for (int i = 0; i < n_nodes; ++i) { - graph_tensor_set.insert(ggml_graph_node(gf, i)); - } - - for (auto& kv : backend_tensor_data_map) { - auto tensor = kv.first; - auto data = kv.second; - if (tensor == nullptr || data == nullptr) { - continue; - } - const char* name = ggml_get_name(tensor); - if (graph_tensor_set.find(tensor) == graph_tensor_set.end()) { - continue; - } - if (tensor->buffer == nullptr) { - LOG_WARN("%s skip backend tensor copy: tensor buffer not set, name='%s', ne=[%lld,%lld,%lld,%lld], type=%s", - get_desc().c_str(), - name != nullptr ? name : "", - (long long)tensor->ne[0], - (long long)tensor->ne[1], - (long long)tensor->ne[2], - (long long)tensor->ne[3], - ggml_type_name(tensor->type)); - continue; - } - - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - if (buf == nullptr) { - LOG_WARN("%s graph exec skip tensor copy: name=%s op=%s reason=buffer_not_set data=%p view_src=%p view_src_buffer=%p", - get_desc().c_str(), - tensor && tensor->name[0] != '\0' ? tensor->name : "", - tensor ? ggml_op_name(tensor->op) : "", - data, - tensor ? tensor->view_src : nullptr, - (tensor && tensor->view_src) ? tensor->view_src->buffer : nullptr); - continue; - } - - ggml_backend_tensor_set(tensor, data, 0, ggml_nbytes(tensor)); - } - - if (clear_after_copy) { - backend_tensor_data_map.clear(); - } -} - -const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_plan(ggml_cgraph* gf) { - GGML_ASSERT(gf != nullptr); - return sd::ggml_graph_cut::resolve_plan(runtime_backend, - gf, - &graph_cut_plan_cache_, - params_tensor_set_, - get_desc().c_str()); -} - -const GGMLRunner::GraphCutPlan& GGMLRunner::resolve_graph_cut_layer_split_plan(ggml_cgraph* gf) { - return resolve_graph_cut_plan(gf); -} - -bool GGMLRunner::assign_graph_cut_layer_split_backends(ggml_cgraph* gf) { - graph_cut_layer_split_node_assignments_.clear(); - if (!graph_cut_layer_split_enabled) { - return true; - } - if (!is_multi_device()) { - LOG_ERROR("%s graph-cut layer split requires multiple runtime backends", get_desc().c_str()); - return false; - } - - const auto& plan = resolve_graph_cut_layer_split_plan(gf); - if (!plan.valid || !plan.has_cuts || plan.segments.size() <= 1) { - auto manager = residency_manager.lock(); - if (manager == nullptr) { - LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); - return false; - } - std::vector graph_params = collect_used_param_tensors(gf); - if (!graph_params.empty() && - !manager->assign_compute_backend(graph_params, runtime_backend)) { - LOG_ERROR("%s graph-cut layer split failed to assign unmarked graph params to %s", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str()); - return false; - } - for (ggml_tensor* param : graph_params) { - if (param != nullptr) { - graph_cut_layer_split_assignments_[param] = runtime_backend; - } - } - const int n_nodes = ggml_graph_n_nodes(gf); - for (int i = 0; i < n_nodes; i++) { - ggml_tensor* node = ggml_graph_node(gf, i); - if (node != nullptr) { - graph_cut_layer_split_node_assignments_[node] = runtime_backend; - } - } - if (!graph_cut_layer_split_primary_notice_logged_) { - LOG_WARN("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str(), - graph_params.size()); - graph_cut_layer_split_primary_notice_logged_ = true; - } else { - LOG_VERBOSE("%s graph-cut layer split: graph has no mark_graph_cut segments; using primary backend %s for %zu graph params", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(runtime_backend).c_str(), - graph_params.size()); - } - return true; - } - - std::vector split_backends; - split_backends.reserve(extra_runtime_backends.size() + 1); - split_backends.push_back(runtime_backend); - for (ggml_backend_t backend : extra_runtime_backends) { - if (backend != nullptr) { - split_backends.push_back(backend); - } - } - - auto manager = residency_manager.lock(); - if (manager == nullptr) { - LOG_ERROR("%s weight manager is not set for graph-cut layer split", get_desc().c_str()); - return false; - } - - sd::GraphCutLayerSplitAssignment assignment; - auto canonicalize_param = [this](ggml_tensor* tensor) { - return canonical_param_tensor(tensor); - }; - if (!sd::partition_graph_cut_layer_split(get_desc().c_str(), - gf, - plan, - split_backends, - graph_cut_layer_split_backend_vram_limits_, - max_graph_vram_bytes, - graph_cut_layer_split_assignments_, - canonicalize_param, - &assignment)) { - return false; - } - - for (size_t i = 0; i < split_backends.size(); i++) { - if (assignment.tensors_by_backend[i].empty()) { - continue; - } - if (!manager->assign_compute_backend(assignment.tensors_by_backend[i], split_backends[i])) { - LOG_ERROR("%s graph-cut layer split failed to assign params to %s", - get_desc().c_str(), - sd::layer_split_backend_device_display_name(split_backends[i]).c_str()); - return false; - } - } - - graph_cut_layer_split_node_assignments_ = std::move(assignment.node_assignments); - sd::log_graph_cut_layer_split_assignment(get_desc().c_str(), split_backends, assignment); - - return true; -} - 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()) { - manager->set_workspace_reclaimer(reinterpret_cast(this), [this]() { - if (!workspace_.release()) { - return false; - } - sync_runtime_residency(); - return true; - }); - } - runner_started_ = true; - return true; + return execution_engine_->start(execution_session_); } void GGMLRunner::runner_end() { GGML_ASSERT(!graph_active_); - if (!runner_started_) { - return; - } - workspace_.release(); + execution_engine_->end(execution_session_); 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()) { - manager->clear_prefetched_params(reinterpret_cast(this)); - std::vector tensors; - for (auto tensor : params_tensor_set_) { - auto* parameter = manager->resolve_param_tensor(const_cast(tensor)); - if (parameter != nullptr) - tensors.push_back(parameter); - } - 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; } GGMLRunner::GGMLRunner(ggml_backend_t backend, std::shared_ptr manager) - : runtime_backend(backend), - cache_(backend), - cut_cache_(backend), - workspace_(backend), + : execution_engine_(manager != nullptr ? manager->execution_engine() : std::make_shared()), + execution_session_(execution_engine_->create_session(backend, manager)), + runtime_backend(backend), residency_manager(manager) { GGML_ASSERT(runtime_backend != nullptr); alloc_params_ctx(); @@ -538,6 +181,7 @@ GGMLRunner::GGMLRunner(ggml_backend_t backend, GGMLRunner::~GGMLRunner() { runner_end(); + execution_engine_->release_session(execution_session_); free_compute_ctx(); free_params_ctx(); } @@ -575,8 +219,7 @@ void GGMLRunner::reset_compute_ctx() { } void GGMLRunner::free_cache_ctx_and_buffer() { - cache_.clear(); - sync_runtime_residency(); + execution_engine_->clear_cache(execution_session_); } void GGMLRunner::set_backend_tensor_data(ggml_tensor* tensor, const void* data) { @@ -610,7 +253,7 @@ void GGMLRunner::cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* if (tensor != nullptr) { ggml_set_output(tensor); } - cache_.stage(name, tensor); + execution_session_.cache().stage(name, tensor); if (graph != nullptr && tensor != nullptr) { // Schedule the cache output here so its source can be reused before graph end. ggml_build_forward_expand(graph, tensor); @@ -646,439 +289,54 @@ std::optional> GGMLRunner::compute(get_graph_cb_t get_graph, GGMLRunner& runner; const bool& success; ~GraphEndGuard() { - if (!runner.workspace_.segment_end()) { - runner.last_compute_status_ = GGML_STATUS_FAILED; + if (!success) { + runner.execution_engine_->discard_graph(runner.execution_session_); } - runner.cache_.graph_end(false); - runner.cut_cache_.clear(); runner.free_compute_ctx(); runner.graph_active_ = false; - if (!success) { - runner.workspace_.release(); - } - runner.sync_runtime_residency(); } } graph_guard{*this, success}; - ggml_cgraph* graph = nullptr; - if (!prepare_compute_graph(get_graph, &graph)) { - return std::nullopt; - } - params_tensor_set_dirty_ = true; - rebuild_params_tensor_set(); - if (auto manager = residency_manager.lock()) { - for (int i = 0; i < sd::ggml_graph_cut::leaf_count(graph); ++i) { - auto* parameter = manager->resolve_param_tensor(sd::ggml_graph_cut::leaf_tensor(graph, i)); - if (parameter != nullptr) - params_tensor_set_.insert(parameter); - } - } - std::optional> output; try { - output = execute_graph(graph, n_threads, no_return, read_outputs); + ggml_cgraph* graph = nullptr; + if (!prepare_compute_graph(get_graph, &graph)) { + return std::nullopt; + } + params_tensor_set_dirty_ = true; + rebuild_params_tensor_set(); + sd::ComputeTask task; + task.description = get_desc(); + task.graph = graph; + task.output = no_return ? nullptr : ggml_get_tensor(compute_ctx, final_result_name.c_str()); + task.weights = collect_used_param_tensors(graph); + task.inputs = backend_tensor_data_map; + task.debug_tensors = debug_tensors; + task.n_threads = n_threads; + task.no_return = no_return; + task.read_outputs = read_outputs; + auto result = execution_engine_->execute(execution_session_, task); + last_compute_status_ = result.status; + success = result.status == GGML_STATUS_SUCCESS; + if (success) { + return std::move(result.output); + } } catch (const std::bad_alloc&) { last_compute_status_ = GGML_STATUS_ALLOC_FAILED; LOG_ERROR("%s graph allocation failed", get_desc().c_str()); - return std::nullopt; } catch (const std::exception& error) { - last_compute_status_ = GGML_STATUS_FAILED; - LOG_ERROR("%s graph execution failed on %s: %s", get_desc().c_str(), - ggml_backend_name(runtime_backend), error.what()); - return std::nullopt; - } - success = output.has_value(); - if (success) { - cache_.graph_end(true); - last_compute_status_ = GGML_STATUS_SUCCESS; + LOG_ERROR("%s graph submission failed: %s", get_desc().c_str(), error.what()); } - return output; + return std::nullopt; } void GGMLRunner::set_graph_cut_layer_split_enabled(bool enabled) { - graph_cut_layer_split_enabled = enabled; - if (!enabled) { - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; - } + execution_session_.set_layer_split_enabled(enabled); } void GGMLRunner::set_graph_cut_layer_split_backend_vram_limits(const std::vector& limits) { - graph_cut_layer_split_backend_vram_limits_ = limits; - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; + execution_session_.set_backend_limits(limits); } void GGMLRunner::set_runtime_backends(const std::vector& backends) { - extra_runtime_backends.clear(); - for (ggml_backend_t backend : backends) { - if (backend == nullptr || backend == runtime_backend) { - continue; - } - if (std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), backend) == - extra_runtime_backends.end()) { - extra_runtime_backends.push_back(backend); - } - } - workspace_.set_extra_backends(extra_runtime_backends); - graph_cut_layer_split_assignments_.clear(); - graph_cut_layer_split_node_assignments_.clear(); - graph_cut_layer_split_primary_notice_logged_ = false; -} - -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, 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; - } - auto placement = graph_cut_layer_split_assignments_.find(tensor); - return placement == graph_cut_layer_split_assignments_.end() ? runtime_backend : placement->second; - }; - auto assign_nodes = [&](ggml_backend_sched_t scheduler, ggml_cgraph* copy) { - pin_multi_device_nodes(scheduler, copy, graph); - }; - return workspace_.measure(graph, external_backend, assign_nodes, !is_multi_device()); -} - -std::vector GGMLRunner::memory_requests( - const std::vector& sizes, - size_t pending_cache_bytes) const { - std::vector requests; - for (const auto& size : sizes) { - const size_t retained = retained_runtime_buffer_bytes(size.backend); - const size_t reusable = workspace_.bytes(size.backend); - const size_t cache_bytes = size.backend == runtime_backend ? pending_cache_bytes : 0; - const size_t pending = add_bytes(size.bytes > reusable ? size.bytes - reusable : 0, cache_bytes); - size_t limit = max_graph_vram_bytes; - if (is_multi_device()) { - size_t index = 0; - if (size.backend != runtime_backend) { - auto position = std::find(extra_runtime_backends.begin(), extra_runtime_backends.end(), size.backend); - index = static_cast(position - extra_runtime_backends.begin()) + 1; - } - if (index < graph_cut_layer_split_backend_vram_limits_.size()) { - limit = graph_cut_layer_split_backend_vram_limits_[index]; - } - } - requests.push_back({size.backend, reinterpret_cast(this), pending, - retained, limit}); - } - return requests; -} - -bool GGMLRunner::fits(const std::vector& requests, - const std::vector& params) const { - auto manager = residency_manager.lock(); - if (manager == nullptr) { - return params.empty(); - } - for (const auto& request : requests) { - if (!manager->fits_compute_backend_capacity(request, params)) { - return false; - } - } - return true; -} - -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); - } - if (workspace_.cpu_backend() != nullptr) { - sd_backend_cpu_set_n_threads(workspace_.cpu_backend(), n_threads); - } - auto scheduler = workspace_.scheduler(); - 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; - LOG_ERROR("%s compute failed: %s", get_desc().c_str(), ggml_status_to_string(status)); - return false; - } - const std::string description = get_desc(); - if (!debug_tensors.empty()) { - std::unordered_set graph_tensors; - const int leaf_count = ggml_graph_cut::leaf_count(graph); - const int node_count = ggml_graph_n_nodes(graph); - graph_tensors.reserve(static_cast(leaf_count + node_count)); - for (int index = 0; index < leaf_count; ++index) { - graph_tensors.insert(ggml_graph_cut::leaf_tensor(graph, index)); - } - for (int index = 0; index < node_count; ++index) { - graph_tensors.insert(ggml_graph_node(graph, index)); - } - - for (const auto& entry : debug_tensors) { - ggml_tensor* tensor = entry.first; - if (tensor == nullptr || graph_tensors.find(tensor) == graph_tensors.end()) { - continue; - } - ggml_backend_buffer_t buffer = - tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer; - if (buffer == nullptr) { - LOG_WARN("%s skip debug tensor '%s': tensor buffer not set", - description.c_str(), - entry.second.c_str()); - continue; - } - if (tensor->type != GGML_TYPE_F32) { - LOG_WARN("%s skip debug tensor '%s': only GGML_TYPE_F32 is supported, got %s", - description.c_str(), - entry.second.c_str(), - ggml_type_name(tensor->type)); - continue; - } - auto debug_tensor = make_sd_tensor_from_ggml(tensor); - print_sd_tensor(debug_tensor, false, entry.second.c_str()); - } - } - - return true; -} - -std::optional> GGMLRunner::execute_graph(ggml_cgraph* graph, int n_threads, bool no_return, const std::function& read_outputs) { - 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); - 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; - } - 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(); - auto& segment = monolithic_plan.segments.back(); - segment.group_name = "graph"; - segment.compute_buffer_size = cached_plan.compute_buffer_size; - segment.internal_node_indices.reserve(ggml_graph_n_nodes(graph)); - segment.input_refs.reserve(ggml_graph_cut::leaf_count(graph)); - for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { - segment.internal_node_indices.push_back(i); - } - for (int i = 0; i < ggml_graph_cut::leaf_count(graph); ++i) { - auto tensor = ggml_graph_cut::leaf_tensor(graph, i); - ggml_graph_cut::Segment::InputRef input; - input.leaf_index = i; - input.type = canonical_param_tensor(tensor) != nullptr - ? ggml_graph_cut::Segment::INPUT_PARAM - : ggml_graph_cut::Segment::INPUT_EXTERNAL; - segment.input_refs.push_back(input); - } - } - const auto& plan = segmented ? cached_plan : monolithic_plan; - const bool segments_changed = plan.segments.size() != logged_segment_count_; - if (segments_changed && (segmented || logged_segment_count_ > 1)) { - LOG_VERBOSE("%s using %zu segment%s", get_desc().c_str(), - plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); - } - SegmentGraphBindings bindings(cut_cache_, plan, graph); - SegmentWeightPipeline weights(manager, runtime_backend, reinterpret_cast(this), - graph, plan, params_tensor_set_, - segmented && manager != nullptr && manager->prefetch_enabled()); - - std::map peak_compute_bytes; - auto track_compute_buffer = [&](ggml_backend_t backend) { - if (backend != nullptr) { - auto& peak = peak_compute_bytes[backend]; - peak = std::max(peak, workspace_.bytes(backend)); - } - }; - std::optional> output = Tensor(); - for (size_t index = 0; index < plan.segments.size(); ++index) { - const auto& segment = plan.segments[index]; - 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); - return std::nullopt; - }; - cut_cache_.prune(segment.live_cut_names); - bindings.reset(segment); - if (!bindings.bind_cached_inputs(segment, get_desc().c_str())) { - return fail_segment("input binding"); - } - ggml_context* segment_context = nullptr; - auto segment_graph = segmented - ? ggml_graph_cut::build_segment_graph(graph, segment, &segment_context) - : graph; - struct SegmentCleanup { - GGMLRunner& runner; - SegmentWeightPipeline& weights; - SegmentGraphBindings& bindings; - ggml_context* context; - ~SegmentCleanup() { - if (!runner.workspace_.segment_end()) { - runner.last_compute_status_ = GGML_STATUS_FAILED; - } - bindings.restore(); - weights.segment_end(); - ggml_free(context); - runner.sync_runtime_residency(); - } - } segment_cleanup{*this, weights, bindings, segment_context}; - - 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); - 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))) { - return true; - } - if (workspace_.release_excess(measurement)) { - sync_runtime_residency(); - requests = memory_requests(measurement.buffers, new_cache_bytes); - } - const bool ready = weights.ensure_segment_capacity(index, requests); - if (!ready && manager != nullptr) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - } - return ready; - }; - if (!weights.segment_start(index, ensure_capacity)) { - return fail_segment("weight preparation"); - } - // Preparing weights can execute LoRA graphs and reclaim an idle workspace. - if (!workspace_.prepare(measurement)) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return fail_segment("workspace preparation"); - } - 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); - })) { - last_compute_status_ = GGML_STATUS_ALLOC_FAILED; - return fail_segment("workspace allocation"); - } - for (const auto& size : measurement.buffers) { - track_compute_buffer(size.backend); - } - if (workspace_.scheduler() != nullptr) { - track_compute_buffer(workspace_.cpu_backend()); - } - if (!ensure_capacity()) { - return fail_segment("allocated capacity check"); - } - copy_data_to_backend_tensor(segment_graph, false); - auto prefetch_requests = memory_requests(measurement.buffers, new_cache_bytes); - if (!prefetch_requests.empty()) { - weights.enqueue_next(index, prefetch_requests.front()); - } - 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, [&]() { weights.submit_next(); })) { - return fail_segment("execution"); - } - auto cache_status = cache_.capture(segment_graph); - if (cache_status == GGML_STATUS_SUCCESS) { - cache_status = cut_cache_.capture(graph, segment, get_desc().c_str()); - } - if (cache_status != GGML_STATUS_SUCCESS) { - last_compute_status_ = cache_status; - return fail_segment("output caching"); - } - sync_runtime_residency(); - if (last) { - if (read_outputs && !read_outputs()) { - return fail_segment("output finalization"); - } - if (!no_return) { - auto result = ggml_get_tensor(compute_ctx, final_result_name.c_str()); - output = read_graph_tensor(result, "output"); - if (!output.has_value()) { - return fail_segment("output readback"); - } - } - } - if (!workspace_.segment_end()) { - last_compute_status_ = GGML_STATUS_FAILED; - return fail_segment("workspace synchronization"); - } - // Final outputs and their callbacks may still be views of consumed cuts. - cut_cache_.prune(segment.future_cut_names); - } - if (segments_changed || peak_compute_bytes != logged_compute_bytes_) { - for (const auto& entry : peak_compute_bytes) { - LOG_VERBOSE("%s compute buffer size: %.2f MB(%s) on %s (peak across %zu segment%s)", - get_desc().c_str(), entry.second / (1024.0 * 1024.0), - sd_backend_is_cpu(entry.first) ? "RAM" : "VRAM", ggml_backend_name(entry.first), - plan.segments.size(), plan.segments.size() == 1 ? "" : "s"); - } - logged_compute_bytes_ = std::move(peak_compute_bytes); - logged_segment_count_ = plan.segments.size(); - } - return output; + execution_session_.set_runtime_backends(backends); } diff --git a/src/core/ggml_runner.h b/src/core/ggml_runner.h index d300be7681..9524824219 100644 --- a/src/core/ggml_runner.h +++ b/src/core/ggml_runner.h @@ -12,9 +12,7 @@ #include #include -#include "core/compute_workspace.h" -#include "core/ggml_graph_cut.h" -#include "core/runner_cache.h" +#include "core/execution_engine.h" #include "core/tensor_ggml.hpp" #include "core/util.h" #include "device_residency_manager.h" @@ -132,43 +130,20 @@ struct GGMLRunnerContext { struct GGMLRunner { private: - std::map logged_compute_bytes_; - size_t logged_segment_count_ = 0; + std::shared_ptr execution_engine_; + sd::ExecutionSession& execution_session_; ggml_status last_compute_status_ = GGML_STATUS_SUCCESS; - 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, 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: typedef std::function get_graph_cb_t; - using GraphCutPlan = sd::ggml_graph_cut::Plan; ggml_backend_t runtime_backend = nullptr; ggml_context* params_ctx = nullptr; - sd::RunnerCache cache_; - sd::GraphCutTensorCache cut_cache_; - sd::ComputeWorkspace workspace_; ggml_context* compute_ctx = nullptr; std::unique_ptr compute_ctx_buffer_; - bool runner_started_ = false; - bool graph_active_ = false; - - size_t max_graph_vram_bytes = 0; - bool graph_cut_layer_split_enabled = false; - std::vector graph_cut_layer_split_backend_vram_limits_; - - std::vector extra_runtime_backends; // borrowed (SDBackendManager-owned) + bool graph_active_ = false; std::shared_ptr weight_adapter = nullptr; std::weak_ptr residency_manager; @@ -193,11 +168,7 @@ struct GGMLRunner { bool circular_x_enabled = false; bool circular_y_enabled = false; - sd::ggml_graph_cut::PlanCache graph_cut_plan_cache_; std::unordered_set params_tensor_set_; - std::unordered_map graph_cut_layer_split_assignments_; - std::unordered_map graph_cut_layer_split_node_assignments_; - bool graph_cut_layer_split_primary_notice_logged_ = false; template static sd::Tensor take_or_empty(std::optional> tensor) { @@ -239,8 +210,6 @@ struct GGMLRunner { std::vector collect_used_param_tensors(ggml_cgraph* gf); - void evict_compute_backend_param_tensors(const std::vector& tensors); - void prepare_build_in_tensor_before(); void prepare_build_in_tensor_after(ggml_cgraph* gf); @@ -252,39 +221,18 @@ struct GGMLRunner { bool prepare_compute_graph(get_graph_cb_t get_graph, ggml_cgraph** gf_out); - ggml_backend_t backend_for_weight(const ggml_tensor* tensor) const; - - // Weightless ops have no scheduler anchor, so pin them to the most recent - // weight device. Views must stay unpinned or cross-device copies can be - // skipped for their consumers. - void pin_multi_device_nodes(ggml_backend_sched_t sched, ggml_cgraph* gf, ggml_cgraph* original_graph = nullptr); - - bool is_multi_device() const { - return !extra_runtime_backends.empty(); - } - size_t reusable_compute_buffer_bytes() const { - return workspace_.bytes(runtime_backend); + return execution_session_.compute_bytes(); } - size_t retained_runtime_buffer_bytes(ggml_backend_t backend = nullptr) const; - - void sync_runtime_residency(); - - std::optional> read_graph_tensor(ggml_tensor* tensor, const char* label); - - void copy_data_to_backend_tensor(ggml_cgraph* gf, bool clear_after_copy = true); - - const GraphCutPlan& resolve_graph_cut_plan(ggml_cgraph* gf); - - const GraphCutPlan& resolve_graph_cut_layer_split_plan(ggml_cgraph* gf); - - bool assign_graph_cut_layer_split_backends(ggml_cgraph* gf); + bool has_cache_tensors() const { + return !execution_session_.cache().empty(); + } public: bool runner_start(); - bool runner_started() const { return runner_started_; } + bool runner_started() const { return execution_session_.started(); } void runner_end(); @@ -335,7 +283,7 @@ struct GGMLRunner { void cache(const std::string name, ggml_tensor* tensor, ggml_cgraph* graph = nullptr); ggml_tensor* get_cache_tensor_by_name(const std::string& name) { - return cache_.get(name); + return execution_session_.cache().get(name); } std::optional> compute(get_graph_cb_t get_graph, @@ -353,7 +301,7 @@ struct GGMLRunner { void set_sage_attention_enabled(bool enabled) { if (sage_attn_enabled != enabled) { free_cache_ctx_and_buffer(); - graph_cut_plan_cache_.graph_cut_plans.clear(); + execution_session_.clear_plan_cache(); sage_attn_enabled = enabled; } } @@ -381,7 +329,7 @@ struct GGMLRunner { } void set_max_graph_vram_bytes(size_t max_vram_bytes) { - max_graph_vram_bytes = max_vram_bytes; + execution_session_.set_max_backend_bytes(max_vram_bytes); } void set_graph_cut_layer_split_enabled(bool enabled); diff --git a/src/device_residency_manager.h b/src/device_residency_manager.h index faf7be27dc..7317410182 100644 --- a/src/device_residency_manager.h +++ b/src/device_residency_manager.h @@ -3,12 +3,17 @@ #include #include +#include #include #include "ggml-backend.h" struct ggml_tensor; +namespace sd { + class ExecutionEngine; +} + enum class WeightPrefetchResult { Scheduled, AlreadyResident, @@ -36,13 +41,45 @@ struct DeviceMemoryRequest { } }; +struct DeviceMemoryCapacity { + size_t required_device_bytes = 0; + size_t required_budget_bytes = 0; + size_t available_device_bytes = SIZE_MAX; + size_t available_budget_bytes = SIZE_MAX; + bool requires_reclamation = false; + bool valid = true; + bool wrapped_device_report = false; + + bool fits(bool after_reclamation = false) const { + return valid && !requires_reclamation && + (after_reclamation || !wrapped_device_report) && + required_device_bytes <= available_device_bytes && + required_budget_bytes <= available_budget_bytes; + } +}; + +struct WeightResidencyBlock { + ggml_backend_buffer_t buffer = nullptr; + ggml_backend_t backend = nullptr; + std::vector tensors; + size_t bytes = 0; + uint64_t last_use_epoch = 0; + bool requires_demotion = false; +}; + struct DeviceResidencyManager { +private: + std::shared_ptr execution_engine_; + +public: virtual ~DeviceResidencyManager() = default; - virtual bool segmented_compute_enabled() const = 0; - virtual bool prefetch_enabled() const = 0; - virtual void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) = 0; - virtual void remove_runtime_owner(uintptr_t owner_id) = 0; + std::shared_ptr execution_engine(); + + virtual bool segmented_compute_enabled() const = 0; + virtual bool prefetch_enabled() const = 0; + virtual void register_runtime_owner(uintptr_t owner_id) = 0; + virtual void remove_runtime_owner(uintptr_t owner_id) = 0; // Capacity requests select their backend's weights; protection spans all backends. virtual bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, const std::vector& required_params) const = 0; @@ -62,6 +99,18 @@ struct DeviceResidencyManager { const std::vector& required_params, const std::vector>& preferred_eviction_order, const std::vector& protected_params) = 0; + virtual DeviceMemoryCapacity inspect_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& params, + bool log_details = false) const { + DeviceMemoryCapacity result; + result.valid = false; + return result; + } + virtual std::vector evictable_compute_backend_blocks( + ggml_backend_t backend, + const std::vector& protected_params) const { return {}; } + virtual bool evict_compute_backend_block(const WeightResidencyBlock& block) { return false; } virtual WeightPrefetchResult prefetch_params( uintptr_t owner_id, const std::vector& tensors) = 0; diff --git a/src/model/diffusion/qwen_image_2_1.hpp b/src/model/diffusion/qwen_image_2_1.hpp index 7cd86fd3e7..00b88f4b38 100644 --- a/src/model/diffusion/qwen_image_2_1.hpp +++ b/src/model/diffusion/qwen_image_2_1.hpp @@ -576,7 +576,7 @@ namespace Qwen { result = run(cache); } if (result.empty() && last_compute_status() == GGML_STATUS_ALLOC_FAILED && - (cache.mode != QwenImage21PrefixCache::Mode::NONE || !cache_.empty())) { + (cache.mode != QwenImage21PrefixCache::Mode::NONE || has_cache_tensors())) { // The failed graph has ended before persistent inputs are released. free_cache_ctx_and_buffer(); prefix_cache_disabled = true; diff --git a/src/model_manager.cpp b/src/model_manager.cpp index 8b6182797f..bfdf2625ec 100644 --- a/src/model_manager.cpp +++ b/src/model_manager.cpp @@ -1,11 +1,9 @@ #include "model_manager.h" #include -#include #include #include #include -#include #include #include @@ -150,7 +148,7 @@ bool ModelManager::set_loras(std::vector loras, SDVersion version) { } if (lora_version_ == version && lora_specs_equal(loras_, loras)) return true; - if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { + if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { return state->pin_count != 0; })) { LOG_ERROR("cannot change LoRA configuration during execution"); @@ -166,7 +164,7 @@ bool ModelManager::set_loras(std::vector loras, SDVersion version) { std::set ModelManager::tensor_names() const { std::set names; for (const auto& state : tensor_states_) { - if (state != nullptr && state->component != ModelComponent::LoRA) { + if (state != nullptr && state->component != ModelComponent::LoRA && model_sources_.count(state->source_file) == 0) { names.insert(state->name); } } @@ -234,7 +232,7 @@ bool ModelManager::register_param_tensors(ModelComponent component, *registered_tensor_size += estimate_tensors_size(tensors); } - const auto scoped_sources = source_file != 0 ? model_loader_.file_tensors(source_file, source_version) : String2TensorStorage{}; + const auto scoped_sources = source_file != 0 ? source_loader(source_file).file_tensors(source_file, source_version) : String2TensorStorage{}; const auto& sources = source_file != 0 ? scoped_sources : model_loader_.get_tensor_storage_map(); std::unordered_set new_tensors; std::vector> new_states; @@ -424,6 +422,20 @@ bool ModelManager::validate_registered_tensors() { return ok; } +bool ModelManager::validate_param_tensors(const std::vector& tensors) { + std::vector states; + if (!resolve_required_tensor_states(tensors, states)) { + return false; + } + for (auto* state : states) { + if (!validate_tensor(*state)) { + return false; + } + state->metadata_validated = true; + } + return true; +} + bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { std::vector need_load; need_load.reserve(states.size()); @@ -500,7 +512,7 @@ bool ModelManager::load_tensors_to_params_backend(const std::vector& states) { - std::map, std::vector> states_by_staging_target; + std::map, std::vector> states_by_staging_target; for (TensorState* state : states) { if (state == nullptr || should_ignore(*state) || is_optional_missing_tensor(state->name)) { continue; @@ -526,12 +538,12 @@ bool ModelManager::stage_tensors_to_compute_backend(const std::vectorcompute_backend); } - states_by_staging_target[{state->compute_backend, staging_buft}].push_back(state); + states_by_staging_target[{state->compute_backend, staging_buft, state->source_file}].push_back(state); } for (const auto& pair : states_by_staging_target) { - ggml_backend_t compute_backend = pair.first.first; - ggml_backend_buffer_type_t staging_buft = pair.first.second; + ggml_backend_t compute_backend = std::get<0>(pair.first); + ggml_backend_buffer_type_t staging_buft = std::get<1>(pair.first); const std::vector& target_states = pair.second; if (target_states.empty()) { continue; @@ -646,7 +658,7 @@ bool ModelManager::apply_loras_to_params(const std::vector& states std::map groups; for (TensorState* state : states) { if (state == nullptr || state->tensor == nullptr || state->component == ModelComponent::LoRA || - state->component == ModelComponent::LatentUpsampler || should_ignore(*state) || is_optional_missing_tensor(state->name)) { + state->component == ModelComponent::LatentUpsampler || model_sources_.count(state->source_file) != 0 || should_ignore(*state) || is_optional_missing_tensor(state->name)) { continue; } if (state->applied_lora_epoch == current_lora_epoch_) { @@ -743,6 +755,9 @@ void ModelManager::reset_lora_applied_params() { } bool ModelManager::should_ignore(const TensorState& state) const { + if (model_sources_.count(state.source_file) != 0) { + return false; + } for (const auto& ignore_prefix : common_ignore_tensors_) { if (starts_with(state.name, ignore_prefix)) { return true; @@ -869,7 +884,7 @@ bool ModelManager::can_mmap_storage(const TensorState& state) const { bool ModelManager::alloc_params_buffers(const std::vector& states, std::vector& created_storage_blocks) { - std::map, std::vector> states_by_buffer_type; + std::map, std::vector> states_by_buffer_type; for (TensorState* state : states) { if (state == nullptr || state->tensor == nullptr) { continue; @@ -878,11 +893,11 @@ bool ModelManager::alloc_params_buffers(const std::vector& states, if (params_buft == nullptr) { return false; } - states_by_buffer_type[{params_buft, static_cast(state->residency_mode)}].push_back(state); + states_by_buffer_type[{params_buft, static_cast(state->residency_mode), state->source_file}].push_back(state); } for (const auto& pair : states_by_buffer_type) { - ggml_backend_buffer_type_t params_buft = pair.first.first; + ggml_backend_buffer_type_t params_buft = std::get<0>(pair.first); const std::vector& states = pair.second; size_t alignment = ggml_backend_buft_get_alignment(params_buft); size_t max_size = ggml_backend_buft_get_max_size(params_buft); @@ -1010,7 +1025,7 @@ bool ModelManager::load_tensor_parts(TensorState& state) { }; const bool success = state.source_file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names, false) - : model_loader_.load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_); + : source_loader(state.source_file).load_file_tensors(state.source_file, state.source_version, callback, names, enable_mmap_); if (!success || loaded != names) { return false; } @@ -1098,7 +1113,7 @@ bool ModelManager::load_tensors(const std::vector& states) { }; const auto file = group.first.first; bool success = file == 0 ? model_loader_.load_tensors(callback, enable_mmap_, &names) - : model_loader_.load_file_tensors(file, group.first.second, callback, names, enable_mmap_); + : source_loader(file).load_file_tensors(file, group.first.second, callback, names, enable_mmap_); if (!success || loaded != names) return false; for (auto& entry : batch) { @@ -1316,7 +1331,7 @@ void ModelManager::erase_params_storage_block(ParamsStorageBlock* block) { void ModelManager::release_all() { clear_all_prefetched_params(); runtime_residencies_.clear(); - workspace_reclaimers_.clear(); + runtime_owners_.clear(); for (auto& state : tensor_states_) { state->pin_count = 0; state->applied_lora_epoch = UINT64_MAX; @@ -1643,11 +1658,17 @@ void ModelManager::evict_compute_backend_params(const std::vector& if (tensors.empty()) { return; } - std::vector required_states; - if (!resolve_required_tensor_states(tensors, required_states)) { + std::unordered_set target_states; + // Session cleanup may follow unregistration, after tensor metadata has been freed. + for (auto* tensor : tensors) { + const auto found = tensor_states_by_tensor_.find(tensor); + if (found != tensor_states_by_tensor_.end()) { + target_states.insert(found->second); + } + } + if (target_states.empty()) { return; } - std::unordered_set target_states(required_states.begin(), required_states.end()); for (const auto& block : compute_staging_blocks_) { const bool intersects = std::any_of( @@ -1743,12 +1764,12 @@ WeightResidencyInfo ModelManager::inspect_compute_backend_params( return info; } -void ModelManager::set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) { - workspace_reclaimers_[owner_id] = std::move(reclaim); +void ModelManager::register_runtime_owner(uintptr_t owner_id) { + runtime_owners_.insert(owner_id); } void ModelManager::remove_runtime_owner(uintptr_t owner_id) { - workspace_reclaimers_.erase(owner_id); + runtime_owners_.erase(owner_id); for (auto it = runtime_residencies_.begin(); it != runtime_residencies_.end();) { if (it->first.first == owner_id) { it = runtime_residencies_.erase(it); @@ -1806,8 +1827,8 @@ 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)); - free_bytes = total_bytes; - result.requires_reclamation = true; + free_bytes = total_bytes; + result.wrapped_device_report = true; } if (total_bytes > 0) { free_bytes = std::min(free_bytes, resident < total_bytes ? total_bytes - resident : 0); @@ -1818,6 +1839,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( if (request.max_backend_bytes > 0) { const size_t resident = add(compute_backend_resident_bytes(request.compute_backend), other_runtime_resident_bytes(request.owner_id, request.compute_backend)); + result.requires_reclamation = result.requires_reclamation || resident > request.max_backend_bytes; result.available_budget_bytes = resident < request.max_backend_bytes ? request.max_backend_bytes - resident : 0; @@ -1842,6 +1864,7 @@ ModelManager::CapacityCheck ModelManager::check_capacity( if (entry.second > 0) { const size_t resident = add(compute_backend_resident_bytes(entry.first), other_runtime_resident_bytes(request.owner_id, entry.first)); + result.requires_reclamation = result.requires_reclamation || resident > entry.second; const size_t available = resident < entry.second ? entry.second - resident : 0; result.available_budget_bytes = std::min(result.available_budget_bytes, available); } @@ -1849,200 +1872,154 @@ ModelManager::CapacityCheck ModelManager::check_capacity( return result; } -bool ModelManager::fits_compute_backend_capacity( +DeviceMemoryCapacity ModelManager::inspect_compute_backend_capacity( const DeviceMemoryRequest& request, - const std::vector& required_params) const { + const std::vector& params, + bool log_details) const { std::vector states; - if (!resolve_required_tensor_states(required_params, states, request.compute_backend)) { - return false; + if (request.compute_backend == nullptr || !resolve_required_tensor_states(params, states, request.compute_backend)) { + DeviceMemoryCapacity result; + result.valid = false; + return result; } - const auto capacity = check_capacity(request, states); - return !capacity.requires_reclamation && capacity.fits(); + return check_capacity(request, states, log_details); } -bool ModelManager::ensure_compute_backend_capacity( - const DeviceMemoryRequest& request, - const std::vector& required_params, - const std::vector>& preferred_eviction_order, - const std::vector& protected_params) { - std::vector required_states; - if (!resolve_required_tensor_states(required_params, required_states, request.compute_backend)) { - return false; - } - - ggml_backend_t compute_backend = request.compute_backend; - if (compute_backend == nullptr) { - LOG_ERROR("model manager cannot reclaim memory for a null compute backend"); - return false; +std::vector ModelManager::evictable_compute_backend_blocks( + ggml_backend_t backend, + const std::vector& protected_params) const { + std::vector result; + if (backend == nullptr) { + return result; } - if (sd_backend_is_cpu(compute_backend)) { - return true; + std::vector protected_list; + if (!resolve_required_tensor_states(protected_params, protected_list)) { + return result; } - - auto fits = [&]() { - const auto capacity = check_capacity(request, required_states); - return !capacity.requires_reclamation && capacity.fits(); - }; - if (fits()) { - return true; + std::unordered_set protected_states(protected_list.begin(), protected_list.end()); + for (const auto& entry : prefetch_blocks_) { + protected_states.insert(entry.second->states.begin(), entry.second->states.end()); } - for (const auto& entry : workspace_reclaimers_) { - if (entry.first != request.owner_id) { - entry.second(); - if (fits()) { - return true; - } + const auto device = ggml_backend_get_device(backend); + const bool cpu = sd_backend_is_cpu(backend); + auto on_device = [&](ggml_backend_buffer_t buffer) { + if (buffer == nullptr) { + return false; } - } - - std::unordered_set protected_states; - std::vector resolved_protected; - if (!resolve_required_tensor_states(protected_params, resolved_protected)) { - return false; - } - protected_states.insert(resolved_protected.begin(), resolved_protected.end()); - for (const auto& entry : prefetch_blocks_) { - if (entry.second != nullptr) { - protected_states.insert(entry.second->states.begin(), entry.second->states.end()); + if (cpu) { + return ggml_backend_buffer_is_host(buffer); } - } - - std::unordered_set eviction_states; - auto add_evictable_state = [&](TensorState* state) { - if (state == nullptr || state->compute_backend != compute_backend || - state->pin_count > 0 || protected_states.find(state) != protected_states.end()) { - return; + if (ggml_backend_buffer_is_host(buffer)) { + return false; } - 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; - if (reloadable && resident) { - eviction_states.insert(state); + const auto type = ggml_backend_buffer_get_type(buffer); + const auto split = split_buffer_devices_.find(type); + if (split != split_buffer_devices_.end()) { + return std::any_of(split->second.begin(), split->second.end(), [&](const auto& entry) { + return ggml_backend_get_device(entry.first) == device; + }); } + return ggml_backend_buft_get_device(type) == device; }; - size_t last_eviction_count = 0; - auto release_eviction_states = [&]() { - 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); - } - // Device memory reports can lag behind releases even when no new candidates are added. - return fits(); - }; - - for (const auto& candidate_params : preferred_eviction_order) { - std::vector candidate_states; - if (!resolve_required_tensor_states(candidate_params, candidate_states)) { - return false; + auto append = [&](ggml_backend_buffer_t buffer, const std::vector& states, bool storage) { + if (!on_device(buffer) || states.empty()) { + return; } - for (TensorState* state : candidate_states) { - add_evictable_state(state); + WeightResidencyBlock candidate; + candidate.buffer = buffer; + candidate.backend = backend; + candidate.bytes = ggml_backend_buffer_get_size(buffer); + for (auto* state : states) { + if (state == nullptr) { + continue; + } + if (state->pin_count != 0 || protected_states.count(state) != 0) { + return; + } + if (storage) { + if (state->staged_to_compute_backend) { + return; + } + if (state->residency_mode != ResidencyMode::Disk) { + if (cpu || state->params_backend == nullptr || sd_backend_is_cpu(state->params_backend) || + state->compute_backend != state->params_backend || + state->sources.empty() || state->applied_lora_epoch != UINT64_MAX || + state->split_buffer_type != nullptr || state->tensor->view_src != nullptr) { + return; + } + candidate.requires_demotion = true; + } + } + candidate.tensors.push_back(state->tensor); + candidate.last_use_epoch = std::max(candidate.last_use_epoch, state->last_use_epoch); } - if (release_eviction_states()) { - return true; + result.push_back(std::move(candidate)); + }; + for (const auto& block : compute_staging_blocks_) { + std::vector states; + states.reserve(block->staged_tensors.size()); + for (const auto& entry : block->staged_tensors) { + states.push_back(entry.first); } + append(block->buffer, states, false); } - - std::vector global_candidates; - global_candidates.reserve(tensor_states_.size()); - for (const auto& state : tensor_states_) { - if (state != nullptr && eviction_states.find(state.get()) == eviction_states.end()) { - global_candidates.push_back(state.get()); - } - } - std::stable_sort(global_candidates.begin(), - global_candidates.end(), - [](const TensorState* lhs, const TensorState* rhs) { - return lhs->last_use_epoch < rhs->last_use_epoch; - }); - for (TensorState* state : global_candidates) { - add_evictable_state(state); - if (release_eviction_states()) { - return true; + for (const auto& block : params_storage_blocks_) { + // Loader-retained mappings cannot be reclaimed by freeing a storage block. + if (block->mmap_tensor_stores.empty()) { + append(block->buffer, block->states, true); } } + return result; +} - // The loader retains mmap buffers after storage blocks are released, so - // demoting mapped states would only force duplicate allocations on reload. - std::unordered_set mapped_states; - for (const auto& block : params_storage_blocks_) { - if (!block->mmap_tensor_stores.empty()) { - mapped_states.insert(block->states.begin(), block->states.end()); +bool ModelManager::evict_compute_backend_block(const WeightResidencyBlock& block) { + const auto candidates = evictable_compute_backend_blocks(block.backend, {}); + const auto found = std::find_if(candidates.begin(), candidates.end(), [&](const auto& candidate) { + return candidate.buffer == block.buffer && candidate.tensors == block.tensors && + candidate.requires_demotion == block.requires_demotion; + }); + if (found == candidates.end()) { + return false; + } + for (auto it = compute_staging_blocks_.begin(); it != compute_staging_blocks_.end(); ++it) { + if ((*it)->buffer == block.buffer) { + free_compute_staging_block(**it); + compute_staging_blocks_.erase(it); + return true; } } - std::unordered_set required_set(required_states.begin(), required_states.end()); - std::vector demoted_states; - auto revert_demotions = [&]() { - for (TensorState* state : demoted_states) { - if (state->loaded_to_params_backend) { - state->residency_mode = ResidencyMode::ParamBackend; + for (auto it = params_storage_blocks_.begin(); it != params_storage_blocks_.end(); ++it) { + if ((*it)->buffer == block.buffer) { + if (block.requires_demotion) { + for (auto* state : (*it)->states) { + if (state != nullptr) { + state->residency_mode = ResidencyMode::Disk; + } + } } - } - }; - auto can_demote_to_disk = [&](TensorState* state) { - return state != nullptr && state->tensor != nullptr && state->tensor->view_src == nullptr && - state->residency_mode == ResidencyMode::ParamBackend && - state->params_backend != nullptr && state->params_backend == state->compute_backend && - state->compute_backend == compute_backend && !sd_backend_is_cpu(state->params_backend) && - state->pin_count == 0 && !state->staged_to_compute_backend && state->loaded_to_params_backend && - protected_states.find(state) == protected_states.end() && - required_set.find(state) == required_set.end() && - mapped_states.find(state) == mapped_states.end() && - !state->sources.empty() && state->applied_lora_epoch == UINT64_MAX && - state->split_buffer_type == nullptr; - }; - for (TensorState* state : global_candidates) { - if (!can_demote_to_disk(state)) { - continue; - } - state->residency_mode = ResidencyMode::Disk; - demoted_states.push_back(state); - add_evictable_state(state); - if (release_eviction_states()) { - revert_demotions(); + free_params_storage_block(**it); + params_storage_blocks_.erase(it); return true; } } - revert_demotions(); + return false; +} - auto capacity = check_capacity(request, required_states); - if (!capacity.requires_reclamation && capacity.fits()) { - return true; +bool ModelManager::fits_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params) const { + std::vector states; + if (!resolve_required_tensor_states(required_params, states, request.compute_backend)) { + return false; } - 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)); - const std::string available_budget = capacity.available_budget_bytes == SIZE_MAX - ? "unlimited" - : sd_format("%.2f MB", capacity.available_budget_bytes / (1024.0 * 1024.0)); - LOG_WARN("model manager cannot make enough memory available on %s: need %.2f MB device / %.2f MB budget, available %s device / %s budget", - ggml_backend_name(compute_backend), - capacity.required_device_bytes / (1024.0 * 1024.0), - capacity.required_budget_bytes / (1024.0 * 1024.0), - available_device.c_str(), - available_budget.c_str()); - return false; + const auto capacity = check_capacity(request, states); + return !capacity.requires_reclamation && capacity.fits(); +} + +bool ModelManager::ensure_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& required_params, + const std::vector>& preferred_eviction_order, + const std::vector& protected_params) { + return execution_engine()->ensure_capacity(*this, request, required_params, preferred_eviction_order, protected_params); } diff --git a/src/model_manager.h b/src/model_manager.h index 81d2546a45..f592a63f11 100644 --- a/src/model_manager.h +++ b/src/model_manager.h @@ -91,6 +91,8 @@ class ModelManager : public DeviceResidencyManager { }; ModelLoader model_loader_; + // Auxiliary models retain their own catalogs and conversion settings. + std::map model_sources_; std::vector> tensor_states_; std::map tensor_states_by_tensor_; mutable std::list resolved_tensor_states_; @@ -101,7 +103,7 @@ class ModelManager : public DeviceResidencyManager { std::map> prefetch_blocks_; std::map prefetch_backends_; std::map, RuntimeResidency> runtime_residencies_; - std::map> workspace_reclaimers_; + std::unordered_set runtime_owners_; bool warned_split_lora_skip_ = false; std::set common_ignore_tensors_; std::vector loras_; @@ -131,6 +133,7 @@ class ModelManager : public DeviceResidencyManager { std::vector& required_states, ggml_backend_t compute_backend = nullptr) const; bool should_ignore(const TensorState& state) const; + const ModelLoader& source_loader(ModelLoader::FileId file) const; bool is_optional_missing_tensor(const std::string& name) const; bool validate_tensor(const TensorState& state) const; @@ -148,17 +151,7 @@ class ModelManager : public DeviceResidencyManager { size_t compute_backend_alloc_size(const std::vector& states, bool missing_only) const; size_t compute_backend_resident_bytes(ggml_backend_t compute_backend) const; - struct CapacityCheck { - size_t required_device_bytes = 0; - size_t required_budget_bytes = 0; - size_t available_device_bytes = SIZE_MAX; - size_t available_budget_bytes = SIZE_MAX; - bool requires_reclamation = false; - bool fits() const { - return required_device_bytes <= available_device_bytes && - required_budget_bytes <= available_budget_bytes; - } - }; + using CapacityCheck = DeviceMemoryCapacity; CapacityCheck check_capacity(const DeviceMemoryRequest& request, const std::vector& states, bool log_details = false) const; @@ -186,6 +179,8 @@ class ModelManager : public DeviceResidencyManager { bool add_file(const std::string& path, const std::string& prefix = "", ModelLoader::FileId* id = nullptr, bool force = false); bool del_file(ModelLoader::FileId id); bool refresh_files(); + ModelLoader::FileId register_model_source(ModelLoader loader); + bool unregister_model_source(ModelLoader::FileId file); ModelLoader::FileVersions source_versions(const std::set& components, const ModelLoader& loader) const; size_t registered_params_size(const std::set& components) const; @@ -267,12 +262,13 @@ class ModelManager : public DeviceResidencyManager { } bool validate_registered_tensors(); + bool validate_param_tensors(const std::vector& tensors); bool load_all_params_eagerly(); bool assign_compute_backend(const std::vector& tensors, ggml_backend_t compute_backend) override; bool prepare_params(const std::vector& tensors) override; - void set_workspace_reclaimer(uintptr_t owner_id, std::function reclaim) override; + void register_runtime_owner(uintptr_t owner_id) override; void remove_runtime_owner(uintptr_t owner_id) override; bool fits_compute_backend_capacity(const DeviceMemoryRequest& request, const std::vector& required_params) const override; @@ -290,6 +286,14 @@ class ModelManager : public DeviceResidencyManager { const std::vector& required_params, const std::vector>& preferred_eviction_order, const std::vector& protected_params) override; + DeviceMemoryCapacity inspect_compute_backend_capacity( + const DeviceMemoryRequest& request, + const std::vector& params, + bool log_details = false) const override; + std::vector evictable_compute_backend_blocks( + ggml_backend_t backend, + const std::vector& protected_params) const override; + bool evict_compute_backend_block(const WeightResidencyBlock& block) override; WeightPrefetchResult prefetch_params( uintptr_t owner_id, const std::vector& tensors) override; diff --git a/src/model_manager_files.cpp b/src/model_manager_files.cpp index adeea95d22..d6453a2108 100644 --- a/src/model_manager_files.cpp +++ b/src/model_manager_files.cpp @@ -74,7 +74,7 @@ void ModelManager::invalidate_sources(const std::unordered_set& st } bool ModelManager::set_loader(ModelLoader loader) { - if (!workspace_reclaimers_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { + if (!runtime_owners_.empty() || std::any_of(tensor_states_.begin(), tensor_states_.end(), [](const auto& state) { return state->pin_count != 0; })) { LOG_ERROR("cannot update model sources during execution"); @@ -87,7 +87,10 @@ bool ModelManager::set_loader(ModelLoader loader) { auto key = std::make_pair(state.source_file, state.source_version); auto found = scoped.find(key); if (found == scoped.end()) - found = scoped.emplace(key, loader.file_tensors(key.first, key.second)).first; + found = scoped.emplace(key, model_sources_.count(key.first) != 0 + ? source_loader(key.first).file_tensors(key.first, key.second) + : loader.file_tensors(key.first, key.second)) + .first; return found->second; }; bool lora_changed = false; @@ -147,7 +150,9 @@ ModelLoader::FileVersions ModelManager::source_versions(const std::setsource_file != 0) { - versions[state->source_file] = loader.file_revision(state->source_file); + versions[state->source_file] = model_sources_.count(state->source_file) != 0 + ? source_loader(state->source_file).file_revision(state->source_file) + : loader.file_revision(state->source_file); continue; } for (const auto& source : find_tensor_sources(*state, sources)) { @@ -167,3 +172,34 @@ size_t ModelManager::registered_params_size(const std::set& comp } return bytes; } + +const ModelLoader& ModelManager::source_loader(ModelLoader::FileId file) const { + auto found = model_sources_.find(file); + return found == model_sources_.end() ? model_loader_ : found->second; +} + +ModelLoader::FileId ModelManager::register_model_source(ModelLoader loader) { + const auto versions = loader.file_versions(); + if (versions.size() != 1) { + LOG_ERROR("a model source must contain exactly one root model file"); + return 0; + } + const auto file = versions.begin()->first; + if (model_loader_.file_revision(file) != 0 || model_sources_.count(file) != 0) { + LOG_ERROR("model source is already registered"); + return 0; + } + loader.set_n_threads(n_threads_); + model_sources_.emplace(file, std::move(loader)); + return file; +} + +bool ModelManager::unregister_model_source(ModelLoader::FileId file) { + if (std::any_of(tensor_states_.begin(), tensor_states_.end(), [file](const auto& state) { + return state->source_file == file; + })) { + LOG_ERROR("cannot unregister a model source while its weights are registered"); + return false; + } + return model_sources_.erase(file) != 0; +} diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 1f7d5e4539..50a4fb7ac1 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -9,10 +9,10 @@ #include "core/rng.hpp" #include "diffusion_engine.h" #include "model/vae/vae.hpp" +#include "pipeline/image_upscale.h" #include "request.h" #include "runtime/denoiser.hpp" #include "runtime/image_preprocess.h" -#include "upscaler.h" namespace sd::pipeline { @@ -751,7 +751,7 @@ namespace sd::pipeline { static sd::Tensor upscale_hires_latent(StableDiffusionGGML* sd, const sd::Tensor& latent, const GenerationRequest& request, - UpscalerGGML* upscaler) { + ImageUpscaleModel* upscaler) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { LOG_ERROR("cancelling hires latent upscale"); return {}; @@ -837,7 +837,7 @@ namespace sd::pipeline { sd->generation_progress.report("upscaling", "hires"); sd::Tensor upscaled_tensor; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { - upscaled_tensor = upscaler->upscale_tensor(decoded); + upscaled_tensor = upscale_image_tensor(*upscaler, decoded, sd->n_threads, request.hires.upscale_tile_size); if (upscaled_tensor.empty()) { LOG_ERROR("hires model upscale failed"); return {}; @@ -1034,22 +1034,18 @@ namespace sd::pipeline { LOG_INFO("hires fix: upscaling to %dx%d", request.hires.target_width, request.hires.target_height); sd->generation_progress.report("upscaling", "hires"); - std::unique_ptr hires_upscaler; + std::unique_ptr hires_upscaler; if (request.hires.upscaler == SD_HIRES_UPSCALER_MODEL) { if (sd->get_cancel_flag() == SD_CANCEL_ALL) { LOG_ERROR("cancelling generation before hires model load"); return false; } 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); - 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, - sd->n_threads)) { + hires_upscaler = load_image_upscale_model(sd->model_manager, + sd->backend_manager, + request.hires.model_path, + sd->max_graph_vram_bytes_for_module(SDBackendModule::UPSCALER)); + if (hires_upscaler == nullptr) { LOG_ERROR("load hires model upscaler failed"); return false; } diff --git a/src/pipeline/image_upscale.cpp b/src/pipeline/image_upscale.cpp new file mode 100644 index 0000000000..1ccb7b553d --- /dev/null +++ b/src/pipeline/image_upscale.cpp @@ -0,0 +1,138 @@ +#include "pipeline/image_upscale.h" + +#include "core/ggml_extend_backend.h" +#include "core/util.h" +#include "model/upscaler/esrgan.hpp" +#include "model/upscaler/mosr.hpp" +#include "model_manager.h" +#include "runtime/tiling.h" + +namespace sd::pipeline { + ImageUpscaleModel::ImageUpscaleModel(const std::shared_ptr& manager) + : manager_(manager) {} + + ImageUpscaleModel::~ImageUpscaleModel() { + if (auto* value = runner()) { + value->runner_end(); + } + if (auto manager = manager_.lock()) { + GGML_ASSERT(manager->unregister_param_tensors(weights_)); + if (source_file_ != 0) { + GGML_ASSERT(manager->unregister_model_source(source_file_)); + } + } + } + + GGMLRunner* ImageUpscaleModel::runner() const { + return mosr_ != nullptr ? static_cast(mosr_.get()) : esrgan_.get(); + } + + int ImageUpscaleModel::scale() const { + if (mosr_ != nullptr) { + return mosr_->config.scale; + } + return esrgan_ != nullptr ? esrgan_->config.scale : 1; + } + + std::unique_ptr load_image_upscale_model( + const std::shared_ptr& manager, + SDBackendManager& backends, + const std::string& path, + size_t max_vram_bytes, + bool conv_direct, + ggml_type weight_type) { + if (manager == nullptr) { + return nullptr; + } + const auto backend = backends.runtime_backend(SDBackendModule::UPSCALER); + const auto params_backend = backends.params_backend(SDBackendModule::UPSCALER); + if (backend == nullptr || params_backend == nullptr) { + LOG_ERROR("failed to initialize upscaler backends"); + return nullptr; + } + ModelLoader loader; + if (!loader.init_from_file_and_convert_name(path, "", VERSION_ESRGAN)) { + LOG_ERROR("init upscaler model loader failed: '%s'", path.c_str()); + return nullptr; + } + loader.set_wtype_override(weight_type); + const auto tensors = loader.get_tensor_storage_map(); + auto model = std::unique_ptr(new ImageUpscaleModel(manager)); + model->source_file_ = manager->register_model_source(std::move(loader)); + if (model->source_file_ == 0) { + return nullptr; + } + std::map weights; + if (tensors.count("gblocks.0.weight") != 0) { + model->mosr_ = std::make_unique(backend, tensors, manager); + if (model->mosr_->mosr_net == nullptr) { + LOG_ERROR("unsupported MoSR weights: '%s'", path.c_str()); + return nullptr; + } + model->mosr_->get_param_tensors(weights); + } else { + model->esrgan_ = std::make_unique(backend, tensors, manager); + if (model->esrgan_->rrdb_net == nullptr) { + LOG_ERROR("unsupported ESRGAN weights: '%s'", path.c_str()); + return nullptr; + } + model->esrgan_->get_param_tensors(weights); + } + for (const auto& entry : weights) { + model->weights_.push_back(entry.second); + } + const auto mode = backends.params_backend_is_disk(SDBackendModule::UPSCALER) + ? ModelManager::ResidencyMode::Disk + : ModelManager::ResidencyMode::ParamBackend; + if (!manager->register_param_tensors(ModelComponent::Upscaler, std::move(weights), mode, + backend, params_backend, nullptr, false, false, nullptr, + model->source_file_, VERSION_ESRGAN) || + !manager->validate_param_tensors(model->weights_)) { + LOG_ERROR("register upscaler weights failed: '%s'", path.c_str()); + return nullptr; + } + model->runner()->set_max_graph_vram_bytes(max_vram_bytes); + model->runner()->set_conv2d_direct_enabled(conv_direct); + LOG_INFO("Upscaler architecture: %s, weight type: %s", model->runner()->get_desc().c_str(), ggml_type_name(weight_type)); + return model; + } + + Tensor upscale_image_tensor(ImageUpscaleModel& model, const Tensor& input, int n_threads, int tile_size) { + if (input.empty() || input.dim() != 4 || (input.shape()[2] != 3 && input.shape()[2] != 4)) { + LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor"); + return {}; + } + struct RunnerEnd { + GGMLRunner* runner; + ~RunnerEnd() { runner->runner_end(); } + } cleanup{model.runner()}; + const bool has_alpha = input.shape()[2] == 4; + auto rgb = has_alpha ? sd::ops::slice(input, 2, 0, 3) : input; + const auto compute = [&](const Tensor& tensor) { + return model.mosr_ != nullptr ? model.mosr_->compute(n_threads, tensor) : model.esrgan_->compute(n_threads, tensor); + }; + const int scale = model.scale(); + Tensor output; + if (tile_size <= 0 || (input.shape()[0] <= tile_size && input.shape()[1] <= tile_size)) { + output = compute(rgb); + } else { + output = process_tiles_2d(rgb, + static_cast(input.shape()[0] * scale), + static_cast(input.shape()[1] * scale), + scale, tile_size, tile_size, 0.25f, false, false, compute); + } + if (output.empty()) { + LOG_ERROR("upscaler compute failed"); + return {}; + } + if (has_alpha) { + auto alpha = sd::ops::slice(input, 2, 3, 4); + auto shape = alpha.shape(); + shape[0] = output.shape()[0]; + shape[1] = output.shape()[1]; + alpha = sd::ops::interpolate(alpha, shape, sd::ops::InterpolateMode::Bilinear); + output = sd::ops::concat(output, alpha, 2); + } + return output; + } +} diff --git a/src/pipeline/image_upscale.h b/src/pipeline/image_upscale.h new file mode 100644 index 0000000000..1de9ca00f4 --- /dev/null +++ b/src/pipeline/image_upscale.h @@ -0,0 +1,57 @@ +#ifndef __SD_PIPELINE_IMAGE_UPSCALE_H__ +#define __SD_PIPELINE_IMAGE_UPSCALE_H__ + +#include +#include +#include +#include + +#include "core/tensor.hpp" +#include "ggml.h" + +class ModelManager; +class SDBackendManager; +struct GGMLRunner; +struct ESRGAN; +struct MoSR; + +namespace sd::pipeline { + class ImageUpscaleModel { + friend std::unique_ptr load_image_upscale_model( + const std::shared_ptr&, + SDBackendManager&, + const std::string&, + size_t, + bool, + ggml_type); + friend Tensor upscale_image_tensor(ImageUpscaleModel&, const Tensor&, int, int); + + std::weak_ptr manager_; + std::unique_ptr esrgan_; + std::unique_ptr mosr_; + std::vector weights_; + uint64_t source_file_ = 0; + + explicit ImageUpscaleModel(const std::shared_ptr& manager); + + public: + ~ImageUpscaleModel(); + GGMLRunner* runner() const; + int scale() const; + }; + + std::unique_ptr load_image_upscale_model( + const std::shared_ptr& manager, + SDBackendManager& backends, + const std::string& path, + size_t max_vram_bytes = 0, + bool conv_direct = false, + ggml_type weight_type = GGML_TYPE_F16); + + Tensor upscale_image_tensor(ImageUpscaleModel& model, + const Tensor& input, + int n_threads, + int tile_size); +} + +#endif // __SD_PIPELINE_IMAGE_UPSCALE_H__ diff --git a/src/upscaler.cpp b/src/upscaler.cpp index c9bbbcb6b2..82aa1f464d 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -2,8 +2,9 @@ #include "core/ggml_extend_backend.h" #include "core/util.h" #include "generation_info.h" +#include "model/upscaler/esrgan.hpp" +#include "model/upscaler/mosr.hpp" #include "model_loader.h" -#include "runtime/tiling.h" #include "stable-diffusion.h" #include @@ -19,188 +20,44 @@ UpscalerGGML::UpscalerGGML(int n_threads, direct(direct), tile_size(tile_size), backend_spec(std::move(backend_spec)), - params_backend_spec(std::move(params_backend_spec)) { -} + params_backend_spec(std::move(params_backend_spec)) {} -UpscalerGGML::~UpscalerGGML() { - // ModelManager holds raw ggml tensor pointers owned by the runner context. - model_manager.reset(); - esrgan_upscaler.reset(); - mosr_upscaler.reset(); -} +UpscalerGGML::~UpscalerGGML() = default; void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { max_graph_vram_bytes = max_vram_bytes; - if (runner() != nullptr) { - runner()->set_max_graph_vram_bytes(max_vram_bytes); - } -} - -GGMLRunner* UpscalerGGML::runner() { - if (mosr_upscaler) { - return mosr_upscaler.get(); + if (model != nullptr) { + model->runner()->set_max_graph_vram_bytes(max_vram_bytes); } - return esrgan_upscaler.get(); } int UpscalerGGML::get_scale() { - if (mosr_upscaler) { - return mosr_upscaler->config.scale; - } - if (esrgan_upscaler) { - return esrgan_upscaler->config.scale; - } - return 1; -} - -sd::Tensor UpscalerGGML::compute(const sd::Tensor& input_tensor) { - if (mosr_upscaler) { - return mosr_upscaler->compute(n_threads, input_tensor); - } - return esrgan_upscaler->compute(n_threads, input_tensor); + return model != nullptr ? model->scale() : 1; } -bool UpscalerGGML::load_from_file(const std::string& esrgan_path, - int n_threads) { +bool UpscalerGGML::load_from_file(const std::string& path, int n_threads) { sd::ParallelScope tensor_scope(&tensor_executor); ggml_log_set(sd_ggml_log_callback, nullptr); - + model.reset(); std::string error; - if (!backend_manager.init(backend_spec.c_str(), - params_backend_spec.c_str(), - /*split_mode_spec=*/nullptr, - &error)) { + if (!backend_manager.init(backend_spec.c_str(), params_backend_spec.c_str(), nullptr, &error)) { LOG_ERROR("upscaler backend config failed: %s", error.c_str()); return false; } - auto backend_for = [&](SDBackendModule module) { - ggml_backend_t module_backend = backend_manager.runtime_backend(module); - if (module_backend == nullptr) { - LOG_ERROR("failed to initialize %s backend", sd_backend_module_name(module)); - } - return module_backend; - }; - auto params_backend_for = [&](SDBackendModule module) { - ggml_backend_t module_backend = backend_manager.params_backend(module); - if (module_backend == nullptr) { - LOG_ERROR("failed to initialize %s params backend", sd_backend_module_name(module)); - } - return module_backend; - }; - auto ensure_backend_pair = [&](SDBackendModule module) { - if (backend_for(module) == nullptr) { - return false; - } - return params_backend_for(module) != nullptr; - }; - if (!ensure_backend_pair(SDBackendModule::UPSCALER)) { - return false; - } - model_manager = std::make_shared(); model_manager->set_n_threads(n_threads); model_manager->set_enable_mmap(false); - - ModelLoader model_loader; - if (!model_loader.init_from_file_and_convert_name(esrgan_path, "", VERSION_ESRGAN)) { - LOG_ERROR("init model loader from file failed: '%s'", esrgan_path.c_str()); - return false; - } - model_loader.set_wtype_override(model_data_type); - LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); - const auto& tensor_storage_map = model_loader.get_tensor_storage_map(); - std::map tensors; - if (tensor_storage_map.count("gblocks.0.weight") > 0) { - mosr_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), - tensor_storage_map, - model_manager); - if (mosr_upscaler->mosr_net == nullptr) { - LOG_ERROR("init mosr model from metadata failed: '%s'", esrgan_path.c_str()); - return false; - } - mosr_upscaler->get_param_tensors(tensors); - } else { - esrgan_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), - tensor_storage_map, - model_manager); - if (esrgan_upscaler->rrdb_net == nullptr) { - LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str()); - return false; - } - esrgan_upscaler->get_param_tensors(tensors); - } - LOG_INFO("Upscaler architecture: %s", runner()->get_desc().c_str()); - runner()->set_max_graph_vram_bytes(max_graph_vram_bytes); - if (direct) { - runner()->set_conv2d_direct_enabled(true); - } - - if (!model_manager->set_loader(model_loader) || - !model_manager->register_param_tensors(ModelComponent::Upscaler, - std::move(tensors), - backend_manager.params_backend_is_disk(SDBackendModule::UPSCALER) ? ModelManager::ResidencyMode::Disk : ModelManager::ResidencyMode::ParamBackend, - backend_for(SDBackendModule::UPSCALER), - params_backend_for(SDBackendModule::UPSCALER)) || - !model_manager->validate_registered_tensors()) { - LOG_ERROR("register upscaler tensors with model manager failed"); - return false; - } - return true; + model = sd::pipeline::load_image_upscale_model(model_manager, backend_manager, path, + max_graph_vram_bytes, direct, model_data_type); + return model != nullptr; } sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { sd::ParallelScope tensor_scope(&tensor_executor); - if (input_tensor.empty() || input_tensor.dim() != 4 || - (input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) { - LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor"); + if (model == nullptr) { return {}; } - - const bool has_alpha = input_tensor.shape()[2] == 4; - sd::Tensor rgb; - if (has_alpha) { - rgb = sd::ops::slice(input_tensor, 2, 0, 3); - } - const sd::Tensor& model_input = has_alpha ? rgb : input_tensor; - sd::Tensor upscaled; - const int scale = get_scale(); - if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { - upscaled = compute(model_input); - } else { - auto on_processing = [&](const sd::Tensor& input_tile) -> sd::Tensor { - auto output_tile = compute(input_tile); - if (output_tile.empty()) { - LOG_ERROR("upscaler compute failed while processing a tile"); - return {}; - } - return output_tile; - }; - - upscaled = process_tiles_2d(model_input, - static_cast(input_tensor.shape()[0] * scale), - static_cast(input_tensor.shape()[1] * scale), - scale, - tile_size, - tile_size, - 0.25f, - false, - false, - on_processing); - } - runner()->runner_end(); - if (upscaled.empty()) { - LOG_ERROR("upscaler compute failed"); - return {}; - } - if (has_alpha) { - auto alpha = sd::ops::slice(input_tensor, 2, 3, 4); - auto alpha_shape = alpha.shape(); - alpha_shape[0] = upscaled.shape()[0]; - alpha_shape[1] = upscaled.shape()[1]; - alpha = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear); - upscaled = sd::ops::concat(upscaled, alpha, 2); - } - return upscaled; + return sd::pipeline::upscale_image_tensor(*model, input_tensor, n_threads, tile_size); } sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor) { diff --git a/src/upscaler.h b/src/upscaler.h index b93065fcc9..f3d09384e6 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -3,9 +3,8 @@ #include "core/ggml_extend_backend.h" #include "core/tensor.hpp" -#include "model/upscaler/esrgan.hpp" -#include "model/upscaler/mosr.hpp" #include "model_manager.h" +#include "pipeline/image_upscale.h" #include "stable-diffusion.h" #include @@ -15,9 +14,7 @@ struct UpscalerGGML { SDBackendManager backend_manager; std::shared_ptr model_manager; ggml_type model_data_type = GGML_TYPE_F16; - std::shared_ptr esrgan_upscaler; - std::shared_ptr mosr_upscaler; - std::string esrgan_path; + std::unique_ptr model; int n_threads; sd::ParallelExecutor tensor_executor; bool direct = false; @@ -36,9 +33,7 @@ struct UpscalerGGML { bool load_from_file(const std::string& esrgan_path, int n_threads); void set_max_graph_vram_bytes(size_t max_vram_bytes); - GGMLRunner* runner(); int get_scale(); - sd::Tensor compute(const sd::Tensor& input_tensor); sd::Tensor upscale_tensor(const sd::Tensor& input_tensor); sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor); }; From 0678d04ae72bc13a0b8fb1392fb8f161bbcb6488 Mon Sep 17 00:00:00 2001 From: leejet Date: Tue, 10 Nov 2026 23:56:14 +0800 Subject: [PATCH 54/57] fix(vae): order Sol attention and FFN writes to reduce VRAM --- src/model/vae/minimax_h3_sol_attention.hpp | 1 + src/model/vae/minimax_h3_vae.hpp | 1 + 2 files changed, 2 insertions(+) diff --git a/src/model/vae/minimax_h3_sol_attention.hpp b/src/model/vae/minimax_h3_sol_attention.hpp index 66d5b5c4fd..701525096a 100644 --- a/src/model/vae/minimax_h3_sol_attention.hpp +++ b/src/model/vae/minimax_h3_sol_attention.hpp @@ -172,6 +172,7 @@ namespace MiniMaxH3VAE { heads, nullptr, false, runner_ctx->flash_attn_enabled); auto indices = ggml_reshape_1d(ctx, batch.output_rows, ggml_nelements(batch.output_rows)); output = ggml_set_rows(ctx, output, ggml_reshape_2d(ctx, attended, dim, indices->ne[0]), indices); + runner_ctx->expand_graph(output); } // Later windows can only read this block and its immediate successor. previous = current; diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index e59eed9444..d4e1c32470 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -430,6 +430,7 @@ namespace MiniMaxH3VAE { auto input = ggml_ext_slice(ctx, x, 1, row, std::min(row + rows_per_chunk, x->ne[1]), false); auto indices = ggml_ext_slice(ctx, sol->row_ids, 0, row, std::min(row + rows_per_chunk, x->ne[1]), false); output = ggml_set_rows(ctx->ggml_ctx, output, finish(input), indices); + ctx->expand_graph(output); } return output; } From 465adf389f252e401257f44efb0ea49da8e4e80b Mon Sep 17 00:00:00 2001 From: leejet Date: Wed, 11 Nov 2026 01:02:14 +0800 Subject: [PATCH 55/57] fix: avoid structured binding capture in MiniMax H3 SOL attention --- src/model/vae/minimax_h3_sol_attention.hpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/model/vae/minimax_h3_sol_attention.hpp b/src/model/vae/minimax_h3_sol_attention.hpp index 701525096a..31df751afa 100644 --- a/src/model/vae/minimax_h3_sol_attention.hpp +++ b/src/model/vae/minimax_h3_sol_attention.hpp @@ -107,7 +107,8 @@ namespace MiniMaxH3VAE { } // Bound the workspace of backends that fall back to unfused attention. constexpr size_t max_windows = 4; - for (const auto& [query_count, windows] : groups) { + for (const auto& group : groups) { + const auto& windows = group.second; for (size_t first_window = 0; first_window < windows.size(); first_window += max_windows) { const auto n = static_cast(std::min(max_windows, windows.size() - first_window)); auto pack_windows = [&](auto member) -> sd::Tensor { From f8417d02e186e33db00c955f7931725ed0e83a16 Mon Sep 17 00:00:00 2001 From: losewayy <1134285612@qq.com> Date: Wed, 11 Nov 2026 01:06:53 +0800 Subject: [PATCH 56/57] perf(cuda): fuse RMSNorm and RoPE in Sol-H3-VAE (#2323) --- docs/ggml_extensions.md | 8 +++--- src/core/cuda/rope.cu | 46 +++++++++++++++++++++++++++----- src/model/vae/minimax_h3_vae.hpp | 14 +++++++++- 3 files changed, 58 insertions(+), 10 deletions(-) diff --git a/docs/ggml_extensions.md b/docs/ggml_extensions.md index 7485938f6a..c643237f56 100644 --- a/docs/ggml_extensions.md +++ b/docs/ggml_extensions.md @@ -215,11 +215,13 @@ 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 +`src/core/cuda/rope.cu` supplies standalone CUDA RoPE, +`RMSNorm -> RoPE` fusion for non-affine normalization, 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 +requires contiguous position encoding and, when present, weights. The non-affine +fusion uses 64-channel F32 inputs and split-half RoPE. 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. diff --git a/src/core/cuda/rope.cu b/src/core/cuda/rope.cu index 2bad5df52a..c012c40424 100644 --- a/src/core/cuda/rope.cu +++ b/src/core/cuda/rope.cu @@ -12,7 +12,7 @@ namespace SD_CUDA_NAMESPACE { v = __fadd_rn(v, cuda_shuffle_xor(v, mask)); return v; } - template + template static __global__ void rms_norm_mul_rope_apply_fused( const float* x, const float* weight, @@ -51,7 +51,7 @@ namespace SD_CUDA_NAMESPACE { #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]); + const float value = affine ? __fmul_rn(__fmul_rn(scale, v[j]), weight[col]) : __fmul_rn(scale, v[j]); float rotated = value; if constexpr (interleaved) { const float other = cuda_shuffle_xor(value, 1, 32); @@ -65,7 +65,7 @@ namespace SD_CUDA_NAMESPACE { } 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 other = affine ? __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]) : __fmul_rn(scale, x_row[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 @@ -178,6 +178,20 @@ namespace SD_CUDA_NAMESPACE { (ggml_nrows(x) + 3) / 4 <= INT_MAX; } + static bool rope_non_affine_fusion_supports(const ggml_tensor* norm, const ggml_tensor* rope) { + if (!norm || norm->op != GGML_OP_RMS_NORM || !rope_is_node(rope) || rope->src[0] != norm || + !rope_supports(nullptr, rope, nullptr) || rope_interleaved(rope)) + return false; + const auto* x = norm->src[0]; + const auto* pe = rope->src[1]; + // The kernel reads contiguous channels and position matrices, with row strides in floats. + return x && x->type == GGML_TYPE_F32 && x->ne[0] == 64 && + x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && + x->nb[2] % sizeof(float) == 0 && x->nb[3] % sizeof(float) == 0 && + ggml_are_same_shape(x, norm) && ggml_is_contiguous(pe) && + ggml_nrows(x) > 0 && (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]; @@ -220,6 +234,9 @@ namespace SD_CUDA_NAMESPACE { if (disabled) return 0; auto* g = const_cast(graph); + if (first + 1 < ggml_graph_n_nodes(g) && + rope_non_affine_fusion_supports(ggml_graph_node(g, first), ggml_graph_node(g, first + 1))) + return 2; 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; @@ -253,6 +270,16 @@ namespace SD_CUDA_NAMESPACE { 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 (rope_non_affine_fusion_supports(rope->src[0], rope)) { + for (int j = i; j >= 0 && j >= i - 64; --j) { + if (ggml_graph_node(graph, j) == rope->src[0]) { + if (j < i) + ggml_ext_graph_delay(graph, j, 1, i + 1); + break; + } + } + continue; + } 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; @@ -296,13 +323,20 @@ namespace SD_CUDA_NAMESPACE { 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* rope = nodes[count <= 2 ? count - 1 : 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])); + const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 2 && overlaps(dst, nodes[1]->src[1])); void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data; CUDAError error; - if (count > 1) { + if (count == 2) { + const auto* pe = rope->src[1]; + const int64_t rows = ggml_nrows(x); + rms_norm_mul_rope_apply_fused<64, false, float, false><<((rows + 3) / 4), dim3(32, 4), 0, stream>>>( + static_cast(x->data), nullptr, static_cast(pe->data), static_cast(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), 1.f); + error = cuda_last_error(); + } else if (count > 2) { error = rms_rope_launch(stream, nodes, output, count > 3, count == 5 ? param_f32(nodes[3], 0) : 1.f); } else { RopeApplyStrides strides; diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index d4e1c32470..881fd01970 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -335,6 +335,18 @@ namespace MiniMaxH3VAE { return {q, k, v}; } + std::array project_sol(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe) { + auto to_qkv = std::dynamic_pointer_cast(blocks["to_qkv"]); + auto qkv_projection = to_qkv->forward(ctx, x); + qkv_projection = ggml_reshape_4d(ctx->ggml_ctx, qkv_projection, 3 * head_dim, num_head, x->ne[1], x->ne[2] * x->ne[3]); + auto qkv = ggml_ext_chunk(ctx, qkv_projection, 3, 0, false); + auto q = ggml_rms_norm(ctx->ggml_ctx, qkv[0], 1e-5f); + auto k = ggml_rms_norm(ctx->ggml_ctx, qkv[1], 1e-5f); + return {Rope::apply_rope(ctx, q, pe, false), + Rope::apply_rope(ctx, k, pe, false), + ggml_ext_cont(ctx, qkv[2])}; + } + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* pe, @@ -343,7 +355,7 @@ namespace MiniMaxH3VAE { auto to_out = std::dynamic_pointer_cast(blocks["to_out"]); if (sol != nullptr) { auto project_rows = [&](ggml_tensor* rows, ggml_tensor* local_pe) { - auto [q, k, v] = project(ctx, norm.forward(ctx, rows), local_pe); + auto [q, k, v] = project_sol(ctx, norm.forward(ctx, rows), local_pe); auto as_rows = [&](ggml_tensor* tensor) { return ggml_reshape_2d(ctx->ggml_ctx, ggml_cont(ctx->ggml_ctx, tensor), dim, rows->ne[1]); }; From 32d516c0197ed5ca31305b8575e14eb1451fad33 Mon Sep 17 00:00:00 2001 From: leejet Date: Wed, 11 Nov 2026 01:24:25 +0800 Subject: [PATCH 57/57] feat: add MiniMax-H3-Image support (#2324) --- README.md | 2 + docs/minimax_h3_image.md | 65 ++++++ examples/common/common.cpp | 6 +- examples/server/recipes/defaults.md | 1 + examples/server/recipes/minimax-h3-image.json | 48 ++++ src/conditioning/conditioner.hpp | 4 +- src/model.h | 7 +- src/model/diffusion/minimax_h3.hpp | 205 +++++++++++++----- src/model/integrations/minimax_h3.cpp | 9 + src/model/vae/minimax_h3_vae.hpp | 111 ++++++---- src/model_loader.cpp | 5 +- src/pipeline/diffusion_engine.cpp | 10 +- src/pipeline/image.cpp | 40 +++- src/pipeline/model_builders.cpp | 8 + src/runtime/denoiser.hpp | 33 +++ 15 files changed, 436 insertions(+), 118 deletions(-) create mode 100644 docs/minimax_h3_image.md create mode 100644 examples/server/recipes/minimax-h3-image.json diff --git a/README.md b/README.md index 0ff9339f38..92103a357d 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,7 @@ API and command-line option may change frequently.*** ## πŸ”₯Important News +* **2026/11/11** πŸš€ stable-diffusion.cpp adds **Day-0 support for MiniMax-H3-Image** * **2026/10/16** πŸš€ stable-diffusion.cpp adds **Day-0 support for MOSS-Image-Base / MOSS-Image-SFT** * **2026/10/16** πŸš€ stable-diffusion.cpp adds **Day-0 support for FLUX.3 Image Klein / Klein Base** * **2026/09/20** πŸš€ stable-diffusion.cpp adds **Day-0 support for Qwen-Image-2.1** @@ -60,6 +61,7 @@ API and command-line option may change frequently.*** - [Z-Image L2P](./docs/z_image_l2p.md) - [MiniT2I](./docs/minit2i.md) - [P-H3-Image](./docs/p_h3_image.md) + - [MiniMax-H3-Image](./docs/minimax_h3_image.md) - [SenseNova U1.5](./docs/sensenova_u1.md) - [Ovis-Image](./docs/ovis_image.md) - [Anima](./docs/anima.md) diff --git a/docs/minimax_h3_image.md b/docs/minimax_h3_image.md new file mode 100644 index 0000000000..0355663d26 --- /dev/null +++ b/docs/minimax_h3_image.md @@ -0,0 +1,65 @@ +# MiniMax-H3-Image + +MiniMax-H3-Image supports text-to-image generation and reference-image editing +with the same checkpoint. It reuses the MiniMax-H3 text encoder and uses a +dedicated image VAE. + +## Model files + +Pass the three MiniMax-H3-Image components separately: + +- `--diffusion-model`: MiniMax-H3-Image diffusion transformer +- `--vae`: MiniMax-H3-Image VAE +- `--llm`: the MiniMax-H3 Qwen3-VL-32B text encoder checkpoint + +Use the MiniMax-H3 text encoder, including its Qwen3-VL vision tower and +DeepStack mergers. An existing MiniMax-H3 text encoder can be reused. If the +vision tower is stored separately, pass it with `--llm_vision`. + +Use a single-file diffusion checkpoint with the six-step PDD output heads +included, such as `minimax_h3_image_pruned_bf16.safetensors` or its INT8 ConvRot +version. Use the MiniMax-H3-Image VAE; the MiniMax-H3 video VAE and P-H3-Image +VAE are incompatible. No audio VAE or P-H3 text connector is needed. + +### Download weights + +- Download minimax_h3_image_pruned + - safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3-Image/tree/main/diffusion_models +- Download qwen3vl_32b_minimax_h3 + - safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3-Image/tree/main/text_encoders + - gguf: https://huggingface.co/leejet/MiniMax-H3-GGUF/tree/main +- Download vae + - safetensors: https://huggingface.co/Comfy-Org/MiniMax-H3-Image/tree/main/vae + +## Text-to-image + +```sh +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\minimax_h3_image_pruned_bf16.safetensors --vae ..\models\vae\minimax_h3_image_vae_fp16.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "A silver tabby kitten holding a sign reading 'sd.cpp', natural daylight" --cfg-scale 1.0 --steps 6 --sampling-method euler --sigmas minimax-h3-image -v -W 1024 -H 1024 --offload-to-cpu --rng cpu --diffusion-fa --vae-tiling +``` + +Use Euler, CFG 1, and `--sigmas minimax-h3-image` for the model's six-step +schedule. The default discrete scheduler also produces this schedule with +`--steps 6`. Prompt expansion is not included. + +## Reference-image editing + +Add `--ref-image` (or `-r`) for reference-image editing: + +```sh +.\bin\Release\sd-cli.exe --diffusion-model ..\models\diffusion_models\minimax_h3_image_pruned_bf16.safetensors --vae ..\models\vae\minimax_h3_image_vae_fp16.safetensors --llm ..\models\text_encoders\qwen3vl_32b_minimax_h3-Q4_K_M.gguf -p "Edit . Put a red scarf on the cat, preserving its face and the background." -r ..\assets\ernie_image\turbo_example.png --cfg-scale 1.0 --steps 6 --sampling-method euler --sigmas minimax-h3-image -v -W 1024 -H 1024 --offload-to-cpu --rng cpu --diffusion-fa --vae-tiling +``` + +`--ref-image` can be repeated. Reference images are numbered in argument order +as ``, ``, and so on, with up to 16 references supported +by the model. Prompts are passed through without an automatic editing prefix. + +Reference images keep their aspect ratio. The text encoder sees a 1024-pixel +short edge, while VAE inputs use a short edge between 1024 and 2048 pixels. +Small references are enlarged and dimensions are rounded to multiples of 32. + +## Shape and runtime notes + +- Width and height are aligned upward to a multiple of 32. +- Start at 1024 Γ— 1024. Larger canvases and more references increase memory use. +- Use quantized weights, `--offload-to-cpu`, and `--vae-tiling` to reduce GPU + memory requirements. diff --git a/examples/common/common.cpp b/examples/common/common.cpp index 70107a396f..102bf9507c 100644 --- a/examples/common/common.cpp +++ b/examples/common/common.cpp @@ -1540,6 +1540,8 @@ ArgOptions SDGenerationParams::get_options() { sigmas_str = "1.0,0.978453,0.95418,0.926626,0.89508,0.845148,0.704534,0.414568,0.0"; } else if (sigmas_str == "p-h3-image") { sigmas_str = "1.0,0.984753,0.967905,0.951688,0.925372,0.881923,0.810083,0.528718,0.0"; + } else if (sigmas_str == "minimax-h3-image") { + sigmas_str = "1.0,0.9675236940383911,0.8875485062599182,0.75,0.5327739119529724,0.23200887441635132,0.0"; } else if (sigmas_str == "hyperflow-h3-flash-4") { sigmas_str = "1.0,0.962575853,0.755781889,0.312497854,0.0"; } else if (sigmas_str == "hyperflow-h3-flash-6") { @@ -1793,11 +1795,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, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.", + "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, minimax-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.", on_sigmas_arg}, {"", "--hires-sigmas", - "custom sigma values for the highres fix second pass, comma-separated (e.g., \"0.85,0.725,0.421875,0.0\") or preset: qwen-image-2.1-turbo, p-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.", + "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, minimax-h3-image, hyperflow-h3-flash-4, hyperflow-h3-flash-6, hyperflow-h3-flash-8.", on_hires_sigmas_arg}, {"", "--skip-layers", diff --git a/examples/server/recipes/defaults.md b/examples/server/recipes/defaults.md index bd0abd5636..0e2b63cb1f 100644 --- a/examples/server/recipes/defaults.md +++ b/examples/server/recipes/defaults.md @@ -70,6 +70,7 @@ For community edit adapters and other optional extensions, these are the base-mo | `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. | +| `minimax-h3-image` | 1 | 6 | 1024 Γ— 1024 | [Runtime guide](../../../docs/minimax_h3_image.md) Dedicated image VAE and fixed PDD sigma schedule. | | `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) | diff --git a/examples/server/recipes/minimax-h3-image.json b/examples/server/recipes/minimax-h3-image.json new file mode 100644 index 0000000000..709cd4e55b --- /dev/null +++ b/examples/server/recipes/minimax-h3-image.json @@ -0,0 +1,48 @@ +{ + "schema_version": 1, + "id": "minimax-h3-image", + "name": "MiniMax H3 Image", + "size_multiple": 32, + "description": "Text-to-image and reference-image editing with the six-step PDD checkpoint. Requires the MiniMax-H3-Image 2D VAE and MiniMax-H3 Qwen3-VL-32B text encoder with its vision tower.", + "model_url": "https://huggingface.co/MiniMaxAI/MiniMax-H3-Image", + "slots": [ + { + "id": "diffusion_model", + "required": true + }, + { + "id": "llm", + "required": true, + "label": "MiniMax-H3 Qwen 3 VL 32B Text Encoder" + }, + { + "id": "llm_vision", + "required": false, + "label": "Vision tower (required when not included in the text encoder)" + }, + { + "id": "vae", + "required": true, + "label": "MiniMax-H3-Image 2D VAE" + } + ], + "defaults": { + "sampling_method": "euler", + "cfg_scale": 1, + "steps": 6, + "width": 1024, + "height": 1024, + "custom_sigmas": [ + 1.0, + 0.9675236940383911, + 0.8875485062599182, + 0.75, + 0.5327739119529724, + 0.23200887441635132, + 0.0 + ] + }, + "features": [ + "ref_images" + ] +} diff --git a/src/conditioning/conditioner.hpp b/src/conditioning/conditioner.hpp index 1534535990..12f4dc9a06 100644 --- a/src/conditioning/conditioner.hpp +++ b/src/conditioning/conditioner.hpp @@ -2589,7 +2589,9 @@ struct LLMEmbedder : public Conditioner { } else if (conditioner_params.ref_images != nullptr) { for (size_t i = 0; i < conditioner_params.ref_images->size(); ++i) { auto resized = resize_for_vision((*conditioner_params.ref_images)[i]); - prompt += ": "; + prompt += version == VERSION_MINIMAX_H3_IMAGE + ? "Picture " + std::to_string(i + 1) + ": " + : ": "; if (!add_vision_outputs(llm->encode_image_outputs(n_threads, resized, false), diff --git a/src/model.h b/src/model.h index 4695f5c74d..2184f39f92 100644 --- a/src/model.h +++ b/src/model.h @@ -71,6 +71,7 @@ enum SDVersion { VERSION_FIBO2, VERSION_CHROMA, VERSION_P_H3_IMAGE, + VERSION_MINIMAX_H3_IMAGE, VERSION_COUNT, }; @@ -167,7 +168,11 @@ static inline bool sd_version_is_ltxav(SDVersion version) { } static inline bool sd_version_is_minimax_h3(SDVersion version) { - return version == VERSION_MINIMAX_H3 || version == VERSION_P_H3_IMAGE; + return version == VERSION_MINIMAX_H3 || version == VERSION_P_H3_IMAGE || version == VERSION_MINIMAX_H3_IMAGE; +} + +static inline bool sd_version_is_minimax_h3_image(SDVersion version) { + return version == VERSION_P_H3_IMAGE || version == VERSION_MINIMAX_H3_IMAGE; } static inline bool sd_version_is_wan(SDVersion version) { diff --git a/src/model/diffusion/minimax_h3.hpp b/src/model/diffusion/minimax_h3.hpp index ea9f9c5aea..2f5cb4f7ce 100644 --- a/src/model/diffusion/minimax_h3.hpp +++ b/src/model/diffusion/minimax_h3.hpp @@ -2,8 +2,10 @@ #define __SD_MODEL_DIFFUSION_MINIMAX_H3_HPP__ #include +#include #include #include +#include #include #include #include @@ -24,6 +26,10 @@ namespace MiniMaxH3 { constexpr float FRAME_RESCALE = 5.f / 3.f; constexpr float VISUAL_COND_TIMESTEP = 0.999f; + inline constexpr std::array pdd_sigmas = { + 1.f, 0.9675236940383911f, 0.8875485062599182f, 0.75f, + 0.5327739119529724f, 0.23200887441635132f, 0.f}; + struct LowRankShape { int64_t q = 0; int64_t k = 0; @@ -52,6 +58,7 @@ namespace MiniMaxH3 { int64_t time_embed_dim = 2688; int64_t rope_inv_freq_len = 16; int64_t adaln_curve_grid = 0; + int64_t pdd_num_inference_steps = 1; int patch_t = 1; int patch_h = 2; int patch_w = 2; @@ -95,6 +102,21 @@ namespace MiniMaxH3 { } if (const auto* weight = find("audio_patch_proj.weight")) { config.audio_latent_channels = weight->logical_ne0(); + } else { + config.audio_latent_channels = 0; + } + if (const auto* weight = find("final_layer.video_out.weight")) { + const int64_t patch_dim = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w; + if (patch_dim <= 0 || weight->ne[1] % patch_dim != 0) { + throw std::runtime_error("invalid MiniMax-H3 output head shape"); + } + config.pdd_num_inference_steps = weight->ne[1] / patch_dim; + if (config.pdd_num_inference_steps != 1 && (config.audio_latent_channels != 0 || config.pdd_num_inference_steps != 6)) { + throw std::runtime_error("unsupported MiniMax-H3 output head bank"); + } + } + if (config.audio_latent_channels == 0 && config.pdd_num_inference_steps != 6) { + throw std::runtime_error("MiniMax-H3-Image requires a checkpoint with its six-step PDD heads included"); } config.num_layers = count_blocks(tensors, prefix + ".blocks."); config.token_refiner_num_layers = count_blocks(tensors, prefix + ".token_refiner.blocks."); @@ -151,14 +173,16 @@ namespace MiniMaxH3 { } LOG_VERBOSE("minimax_h3: layers=%" PRId64 ", hidden=%" PRId64 ", heads=%" PRId64 - ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 ", two_time=%d", + ", head_dim=%" PRId64 ", ffn=%" PRId64 ", adaln_curve=%" PRId64 ", two_time=%d, audio_latent_channels=%" PRId64 ", pdd_num_inference_steps=%" PRId64, config.num_layers, config.hidden_size, config.num_attention_heads, config.attention_head_dim, config.ffn_hidden_size, config.adaln_curve_grid, - config.two_time); + config.two_time, + config.audio_latent_channels, + config.pdd_num_inference_steps); return config; } }; @@ -624,6 +648,27 @@ namespace MiniMaxH3 { } }; + static std::array pdd_head_weights(float sigma, float sigma_next) { + auto nearest = [](float value) { + size_t index = 0; + for (size_t i = 1; i < pdd_sigmas.size(); ++i) { + if (std::abs(pdd_sigmas[i] - value) < std::abs(pdd_sigmas[index] - value)) { + index = i; + } + } + return index; + }; + const size_t start = std::min(nearest(sigma), pdd_sigmas.size() - 2); + const size_t stop = std::max(nearest(sigma_next), start + 1); + std::array weights{}; + // Comfy checkpoints keep head zero as the base and store offsets in the remaining heads. + weights[0] = 1.f; + for (size_t i = std::max(start, 1); i < stop; ++i) { + weights[i] = (pdd_sigmas[i] - pdd_sigmas[i + 1]) / (pdd_sigmas[start] - pdd_sigmas[stop]); + } + return weights; + } + struct FinalLayer : public GGMLBlock { Config config; @@ -637,19 +682,23 @@ namespace MiniMaxH3 { 1, !config.uses_adaln_curves(), config.uses_adaln_curves()); - blocks["video_out"] = std::make_shared(config.hidden_size, video_dim, true, true); - blocks["audio_out"] = std::make_shared(config.hidden_size, config.audio_latent_channels, true, true); + blocks["video_out"] = + std::make_shared(config.hidden_size, video_dim * config.pdd_num_inference_steps, true, true); + if (config.audio_latent_channels > 0) { + blocks["audio_out"] = std::make_shared(config.hidden_size, config.audio_latent_channels, true, true); + } } std::pair forward(GGMLRunnerContext* ctx, ggml_tensor* x, ggml_tensor* t_emb, const TokenModulationSpan& video, - const TokenModulationSpan& audio) { + const TokenModulationSpan& audio, + float sigma, + float sigma_next) { auto norm = std::dynamic_pointer_cast(blocks["norm"]); auto adaln = std::dynamic_pointer_cast(blocks["adaln_proj"]); auto video_out = std::dynamic_pointer_cast(blocks["video_out"]); - auto audio_out = std::dynamic_pointer_cast(blocks["audio_out"]); auto mods = adaln->forward(ctx, t_emb); auto apply = [&](const TokenModulationSpan& segment) { auto row = modulation_row(ctx, mods, config.hidden_size, 2, 1, segment.modulation_row); @@ -658,8 +707,25 @@ namespace MiniMaxH3 { 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))}; + auto prediction = video_out->forward(ctx, apply(video)); + if (config.pdd_num_inference_steps > 1) { + GGML_ASSERT(std::isfinite(sigma_next)); + auto weights = pdd_head_weights(sigma, sigma_next); + const int64_t dim = prediction->ne[0] / config.pdd_num_inference_steps; + auto blended = ggml_ext_scale(ctx, ggml_ext_slice(ctx, prediction, 0, 0, dim), weights[0]); + for (int64_t i = 1; i < config.pdd_num_inference_steps; ++i) { + if (weights[i] != 0.f) { + auto offset = ggml_ext_slice(ctx, prediction, 0, i * dim, (i + 1) * dim); + blended = ggml_add(ctx->ggml_ctx, blended, ggml_ext_scale(ctx, offset, weights[i])); + } + } + prediction = blended; + } + ggml_tensor* audio_prediction = nullptr; + if (config.audio_latent_channels > 0) { + audio_prediction = std::dynamic_pointer_cast(blocks["audio_out"])->forward(ctx, apply(audio)); + } + return {prediction, audio_prediction}; } }; @@ -670,8 +736,10 @@ namespace MiniMaxH3 { : config(config) { int64_t video_dim = config.video_latent_channels * config.patch_t * config.patch_h * config.patch_w; blocks["video_patch_proj"] = std::make_shared(video_dim, config.hidden_size, true, true); - blocks["audio_patch_proj"] = std::make_shared(config.audio_latent_channels, config.hidden_size, true, true); - blocks["condition_proj"] = std::make_shared(config.text_dim, config.hidden_size, true); + if (config.audio_latent_channels > 0) { + blocks["audio_patch_proj"] = std::make_shared(config.audio_latent_channels, config.hidden_size, true, true); + } + blocks["condition_proj"] = std::make_shared(config.text_dim, config.hidden_size, true); if (!config.uses_adaln_curves()) { blocks["time_embedder"] = std::make_shared(config.timestep_input_dim, config.time_embed_hidden_size, @@ -788,9 +856,10 @@ namespace MiniMaxH3 { const std::vector& segments, const std::vector& sequence_segments, const TokenModulationSpan& video_segment, - const TokenModulationSpan& audio_segment) { + const TokenModulationSpan& audio_segment, + float sigma, + float sigma_next) { auto video_proj = std::dynamic_pointer_cast(blocks["video_patch_proj"]); - auto audio_proj = std::dynamic_pointer_cast(blocks["audio_patch_proj"]); std::vector> condition_video_ranges; ggml_tensor* video_rows = nullptr; @@ -836,11 +905,15 @@ namespace MiniMaxH3 { audio_offset += rows->ne[1]; audio_rows = audio_rows == nullptr ? rows : ggml_concat(ctx->ggml_ctx, audio_rows, rows, 1); } - audio = pack_audio_rows(audio); - std::pair target_audio_range = {audio_offset, audio_offset + audio->ne[1]}; - audio_rows = audio_rows == nullptr ? audio : ggml_concat(ctx->ggml_ctx, audio_rows, audio, 1); - auto audio_embeds = audio_proj->forward(ctx, audio_rows); - context = refine_context(ctx, context); + std::pair target_audio_range{}; + ggml_tensor* audio_embeds = nullptr; + if (audio != nullptr) { + audio = pack_audio_rows(audio); + target_audio_range = {audio_offset, audio_offset + audio->ne[1]}; + audio_rows = audio_rows == nullptr ? audio : ggml_concat(ctx->ggml_ctx, audio_rows, audio, 1); + audio_embeds = std::dynamic_pointer_cast(blocks["audio_patch_proj"])->forward(ctx, audio_rows); + } + context = refine_context(ctx, context); ggml_tensor* h = nullptr; auto append = [&](ggml_tensor* value) { @@ -891,7 +964,7 @@ namespace MiniMaxH3 { } auto final_layer = std::dynamic_pointer_cast(blocks["final_layer"]); - auto output = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment); + auto output = final_layer->forward(ctx, h, t_emb, video_segment, audio_segment, sigma, sigma_next); auto video_out = DiT::unpatchify_3d(ctx, output.first, video->ne[2] / config.patch_t, @@ -901,15 +974,18 @@ namespace MiniMaxH3 { config.patch_h, config.patch_w, true); - auto audio_out = ggml_reshape_4d(ctx->ggml_ctx, - output.second, - config.audio_latent_channels, - audio->ne[1] / 2, - 2, - audio->ne[2]); - audio_out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, audio_out, 1, 2, 0, 3)); video_out = ggml_ext_scale(ctx, video_out, -1.f); - audio_out = ggml_ext_scale(ctx, audio_out, -1.f); + if (audio == nullptr) { + return {video_out, nullptr}; + } + auto audio_out = ggml_reshape_4d(ctx->ggml_ctx, + output.second, + config.audio_latent_channels, + audio->ne[1] / 2, + 2, + audio->ne[2]); + audio_out = ggml_cont(ctx->ggml_ctx, ggml_ext_torch_permute(ctx, audio_out, 1, 2, 0, 3)); + audio_out = ggml_ext_scale(ctx, audio_out, -1.f); return {video_out, audio_out}; } }; @@ -995,7 +1071,7 @@ namespace MiniMaxH3 { int64_t run_start = 0; int current_tag = text_tags.empty() ? 1 : text_tags[0]; for (int64_t i = 1; i <= text_len; ++i) { - int tag = i < text_len && !text_tags.empty() ? text_tags[i] : -1; + int tag = i < text_len ? (text_tags.empty() ? 1 : text_tags[i]) : -1; if (i == text_len || tag != current_tag) { layout.segments.push_back({run_start, i, @@ -1121,16 +1197,18 @@ namespace MiniMaxH3 { } } - int64_t audio_start = row; - append_audio_positions(audio_t, cursor, w_axis.front(), w_axis.back()); - layout.audio_segment = {audio_start, row + audio_t * 2, audio_time_row}; - layout.sequence_segments.push_back({audio_start, - row + audio_t * 2, - SequenceKind::TARGET_AUDIO}); - layout.segments.push_back({audio_start, - row + audio_t * 2, - audio_time_row * 3 + 2}); - row += audio_t * 2; + if (audio_t > 0) { + int64_t audio_start = row; + append_audio_positions(audio_t, cursor, w_axis.front(), w_axis.back()); + layout.audio_segment = {audio_start, row + audio_t * 2, audio_time_row}; + layout.sequence_segments.push_back({audio_start, + row + audio_t * 2, + SequenceKind::TARGET_AUDIO}); + layout.segments.push_back({audio_start, + row + audio_t * 2, + audio_time_row * 3 + 2}); + row += audio_t * 2; + } int64_t video_start = row; for (int64_t t = 0; t < latent_t; ++t) { @@ -1232,26 +1310,29 @@ namespace MiniMaxH3 { float audio_shift, float video_sigma_next) { GGML_ASSERT(packed.dim() == 4 || packed.dim() == 5); - GGML_ASSERT(audio_length > 0); + GGML_ASSERT(audio_length > 0 || config.audio_latent_channels == 0); GGML_ASSERT(!context_tensor.empty()); 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); + 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); + ggml_tensor* audio_carrier = nullptr; + if (config.audio_latent_channels > 0) { + 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()); for (const auto& condition : condition_videos) { @@ -1282,7 +1363,7 @@ namespace MiniMaxH3 { // The sampler carries c_a = (sigma_v / sigma_a) * x_a so the packed // latent follows one sigma schedule. Restore x_a for the H3 network. auto audio_ctx = get_context(); - auto audio = ggml_ext_scale(&audio_ctx, audio_carrier, sigma_a / sigma_v); + auto audio = audio_carrier == nullptr ? nullptr : ggml_ext_scale(&audio_ctx, audio_carrier, sigma_a / sigma_v); auto layout = build_layout(context_tensor.shape()[1], packed.shape()[2], packed.shape()[1], @@ -1370,14 +1451,18 @@ namespace MiniMaxH3 { layout.segments, layout.sequence_segments, layout.video_segment, - layout.audio_segment); + layout.audio_segment, + sigma_v, + video_sigma_next); // Convert the model's audio velocity to d(c_a) / d(sigma_v). - output.second = ggml_add(compute_ctx, - ggml_ext_scale(&runner_ctx, audio, 1.f - audio_scale), - ggml_ext_scale(&runner_ctx, - output.second, - 1.f + (audio_scale - 1.f) * sigma_a)); - auto merged = merge_av_latents(&runner_ctx, output.first, output.second); + if (output.second != nullptr) { + output.second = ggml_add(compute_ctx, + ggml_ext_scale(&runner_ctx, audio, 1.f - audio_scale), + ggml_ext_scale(&runner_ctx, + output.second, + 1.f + (audio_scale - 1.f) * sigma_a)); + } + auto merged = output.second == nullptr ? output.first : merge_av_latents(&runner_ctx, output.first, output.second); ggml_build_forward_expand(graph, merged); return graph; } diff --git a/src/model/integrations/minimax_h3.cpp b/src/model/integrations/minimax_h3.cpp index 7624a6cc13..c4a7ae81ad 100644 --- a/src/model/integrations/minimax_h3.cpp +++ b/src/model/integrations/minimax_h3.cpp @@ -41,6 +41,15 @@ const std::vector& minimax_h3_definitions() { definition.build = build_minimax_h3; result.push_back(std::move(definition)); } + { + ModelDefinition definition = result.front(); + definition.version = VERSION_MINIMAX_H3_IMAGE; + definition.name = "MiniMax-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)); + } { ModelDefinition definition = result.front(); definition.version = VERSION_P_H3_IMAGE; diff --git a/src/model/vae/minimax_h3_vae.hpp b/src/model/vae/minimax_h3_vae.hpp index 881fd01970..c43c2c7152 100644 --- a/src/model/vae/minimax_h3_vae.hpp +++ b/src/model/vae/minimax_h3_vae.hpp @@ -26,14 +26,19 @@ namespace MiniMaxH3VAE { } struct MiniMaxH3VAEConfig { - int num_layers = 36; - int patch_size_t = 4; - bool sol = false; + int num_layers = 36; + int patch_size_t = 4; + int hidden_size = 2048; + bool sol = false; + bool vit_decoder_2d = false; static MiniMaxH3VAEConfig detect_from_weights(const String2TensorStorage& tensors, const std::string& prefix) { MiniMaxH3VAEConfig config; config.sol = tensors.count(prefix + ".decoder.blocks.0.scale1") != 0; auto proj = tensors.find(prefix + ".decoder.proj_out.weight"); + if (const auto embedder = tensors.find(prefix + ".decoder.x_embedder.weight"); embedder != tensors.end()) { + config.hidden_size = static_cast(embedder->second.ne[1]); + } 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) { @@ -50,11 +55,13 @@ namespace MiniMaxH3VAE { config.num_layers = layers; } } + config.vit_decoder_2d = config.patch_size_t == 1 && config.hidden_size == 1024; if (config.sol || proj != tensors.end()) { - LOG_VERBOSE("minimax_h3_vae: decoder=%s, layers=%d, patch_size_t=%d", + LOG_VERBOSE("minimax_h3_vae: decoder=%s, layers=%d, patch_size_t=%d, hidden_size=%d", config.sol ? "sol" : "original", config.num_layers, - config.patch_size_t); + config.patch_size_t, + config.hidden_size); } return config; } @@ -287,11 +294,12 @@ namespace MiniMaxH3VAE { } struct DecoderAttention : public GGMLBlock { - static constexpr int num_head = 32; static constexpr int head_dim = 64; - static constexpr int dim = num_head * head_dim; + int dim; + int num_head; - DecoderAttention() { + explicit DecoderAttention(int dim = 2048) + : dim(dim), num_head(dim / head_dim) { blocks["to_qkv"] = std::make_shared(dim, dim * 3, true); blocks["to_out"] = std::make_shared(dim, dim, true); } @@ -366,26 +374,23 @@ namespace MiniMaxH3VAE { }; return to_out->forward(ctx, sol->forward(ctx, x, pe, num_head, project_rows)); } - auto [q, k, v] = project(ctx, norm.forward(ctx, x), pe); - auto out = ggml_ext_attention_ext(ctx, - q, - k, - v, - num_head, - nullptr, - true, - ctx->flash_attn_enabled); + auto [q, k, v] = project(ctx, norm.forward(ctx, x), pe); + auto out = ggml_ext_attention_ext(ctx, + q, + k, + v, + num_head, + nullptr, + true, + ctx->flash_attn_enabled); return to_out->forward(ctx, out); } }; struct DecoderFeedForward : public GGMLBlock { - static constexpr int dim = 2048; - static constexpr int kInnerDim = dim * 4; - - DecoderFeedForward() { - blocks["w1"] = std::make_shared(dim, kInnerDim * 2, true); - blocks["w2"] = std::make_shared(kInnerDim, dim, true); + explicit DecoderFeedForward(int dim = 2048) { + blocks["w1"] = std::make_shared(dim, dim * 8, true); + blocks["w2"] = std::make_shared(dim * 4, dim, true); } ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { @@ -397,13 +402,14 @@ namespace MiniMaxH3VAE { }; struct DecoderBlock : public GGMLBlock { - static constexpr int dim = 2048; + int dim; - DecoderBlock() { + explicit DecoderBlock(int dim = 2048) + : dim(dim) { blocks["norm1"] = std::make_shared(dim, 1e-5f); - blocks["attn"] = std::make_shared(); + blocks["attn"] = std::make_shared(dim); blocks["norm2"] = std::make_shared(dim, 1e-5f); - blocks["ff"] = std::make_shared(); + blocks["ff"] = std::make_shared(dim); } void init_params(ggml_context* ctx, @@ -420,10 +426,10 @@ namespace MiniMaxH3VAE { ggml_tensor* pe, const SolAttentionPlan::Inputs* sol, int64_t rows_per_chunk) { - auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); - auto attn = std::dynamic_pointer_cast(blocks["attn"]); - auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); - auto ff = std::dynamic_pointer_cast(blocks["ff"]); + auto norm1 = std::dynamic_pointer_cast(blocks["norm1"]); + auto attn = std::dynamic_pointer_cast(blocks["attn"]); + auto norm2 = std::dynamic_pointer_cast(blocks["norm2"]); + auto ff = std::dynamic_pointer_cast(blocks["ff"]); x = ggml_add(ctx->ggml_ctx, x, ggml_mul(ctx->ggml_ctx, @@ -449,7 +455,7 @@ namespace MiniMaxH3VAE { }; struct Decoder : public GGMLBlock { - static constexpr int dim = 2048; + int dim; int num_layers; static constexpr int num_register_tokens = 4; static constexpr int patch_size = 16; @@ -457,11 +463,11 @@ namespace MiniMaxH3VAE { bool sol; explicit Decoder(const MiniMaxH3VAEConfig& config) - : num_layers(config.num_layers), patch_size_t(config.patch_size_t), sol(config.sol) { + : dim(config.hidden_size), num_layers(config.num_layers), patch_size_t(config.patch_size_t), sol(config.sol) { blocks["x_embedder"] = std::make_shared(24, dim, true); for (int i = 0; i < num_layers; ++i) { blocks[(sol ? "blocks." : "transformer_blocks.") + std::to_string(i)] = - std::make_shared(); + std::make_shared(dim); } blocks["norm_out"] = std::make_shared(dim, 1e-5f, true, true); blocks["proj_out"] = std::make_shared(dim, @@ -574,15 +580,19 @@ namespace MiniMaxH3VAE { struct MiniMaxH3VideoVAE : public GGMLBlock { bool single_frame; + bool vit_decoder_2d; explicit MiniMaxH3VideoVAE(const MiniMaxH3VAEConfig& config) - : single_frame(config.patch_size_t == 1) { + : single_frame(config.patch_size_t == 1), vit_decoder_2d(config.vit_decoder_2d) { blocks["encoder"] = std::make_shared(); blocks["quant_conv"] = std::make_shared(48, 48, std::tuple{1, 1, 1}); blocks["post_quant_conv"] = std::make_shared(24, 24, std::tuple{1, 1, 1}); - blocks["decoder"] = std::make_shared(config); + if (vit_decoder_2d) { + blocks["post_quant_conv"] = std::make_shared(24, 24, std::pair{1, 1}); + } + blocks["decoder"] = std::make_shared(config); } ggml_tensor* encode(GGMLRunnerContext* ctx, @@ -605,9 +615,17 @@ namespace MiniMaxH3VAE { const std::array& sol_inputs, ggml_tensor* pixel_mean, ggml_tensor* pixel_std) { - auto post_quant = std::dynamic_pointer_cast(blocks["post_quant_conv"]); + auto post_quant = std::dynamic_pointer_cast(blocks["post_quant_conv"]); auto decoder = std::dynamic_pointer_cast(blocks["decoder"]); - return decoder->forward(ctx, post_quant->forward(ctx, latent), pe, shifted_pe, sol_inputs, pixel_mean, pixel_std); + if (vit_decoder_2d) { + GGML_ASSERT(latent->ne[2] == 1); + latent = ggml_reshape_4d(ctx->ggml_ctx, latent, latent->ne[0], latent->ne[1], 24, latent->ne[3] / 24); + } + latent = post_quant->forward(ctx, latent); + if (vit_decoder_2d) { + latent = ggml_reshape_4d(ctx->ggml_ctx, latent, latent->ne[0], latent->ne[1], 1, latent->ne[2] * latent->ne[3]); + } + return decoder->forward(ctx, latent, pe, shifted_pe, sol_inputs, pixel_mean, pixel_std); } }; @@ -651,6 +669,9 @@ namespace MiniMaxH3VAE { if (config.sol) { return "sol_h3_vae"; } + if (config.vit_decoder_2d) { + return "minimax_h3_image_vae"; + } return config.patch_size_t == 1 ? "p_h3_image_vae" : "minimax_h3_video_vae"; } @@ -759,7 +780,7 @@ namespace MiniMaxH3VAE { auto input = ensure_video_shape(x); auto tiling = resolve_original_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"); + LOG_ERROR("H3 image VAE encodes single-frame images only"); return {}; } if (input.shape()[2] == 1) { @@ -865,7 +886,7 @@ namespace MiniMaxH3VAE { } 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"); + LOG_ERROR("H3 image VAE decodes single-frame latents only"); return {}; } if (input.shape()[2] == 1) { @@ -966,7 +987,7 @@ namespace MiniMaxH3VAE { std::vector> ids; const int64_t global_rows = config.sol ? ((num_frames + sol_block_size - 1) / sol_block_size) * sol_global_rows : 5; ids.reserve(static_cast(width * height * num_frames + global_rows)); - constexpr float two_pi = 6.28318530717958647692f; + constexpr float two_pi = 6.28318530717958647692f; const int64_t spatial_shift = shifted ? sol_spatial_shift : 0; const int64_t time_shift = shifted ? sol_temporal_shift : 0; for (int64_t t = 0; t < num_frames; ++t) { @@ -979,20 +1000,20 @@ namespace MiniMaxH3VAE { for (int64_t w = 0; w < width; ++w) { const float x = config.sol ? ((w + spatial_shift) % sol_window_size + 0.5f) / sol_window_size : (w + 0.5f) / width; float pw = (2.f * x - 1.f) * two_pi; - ids.push_back({pt, ph, pw}); + ids.push_back(config.vit_decoder_2d ? std::vector{ph, pw} : std::vector{pt, ph, pw}); } } } for (int64_t i = 0; i < global_rows; ++i) { - ids.push_back({0.f, 0.f, 0.f}); + ids.push_back(std::vector(config.vit_decoder_2d ? 2 : 3, 0.f)); } auto values = Rope::embed_nd(ids, 1, 100.f, - std::vector{16, 16, 16}); + config.vit_decoder_2d ? std::vector{16, 16} : std::vector{16, 16, 16}); return sd::Tensor({2, 2, - 24, + config.vit_decoder_2d ? 16 : 24, static_cast(ids.size())}, std::move(values)); } diff --git a/src/model_loader.cpp b/src/model_loader.cpp index 73025a4075..5987146c20 100644 --- a/src/model_loader.cpp +++ b/src/model_loader.cpp @@ -577,7 +577,10 @@ SDVersion ModelLoader::get_sd_version() const { return VERSION_LTXAV; } 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()) { + tensor_storage_map.find("model.diffusion_model.condition_proj.weight") != tensor_storage_map.end()) { + if (tensor_storage_map.count("model.diffusion_model.audio_patch_proj.weight") == 0) { + return VERSION_MINIMAX_H3_IMAGE; + } // 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()) { diff --git a/src/pipeline/diffusion_engine.cpp b/src/pipeline/diffusion_engine.cpp index b49ef28c4c..5666680e30 100644 --- a/src/pipeline/diffusion_engine.cpp +++ b/src/pipeline/diffusion_engine.cpp @@ -41,6 +41,7 @@ #include "model/audio/wav2vec2.hpp" #include "model/diffusion/animatediff.hpp" #include "model/diffusion/control.hpp" +#include "model/diffusion/minimax_h3.hpp" #include "model/diffusion/model.hpp" #include "model/diffusion/z_image.hpp" #include "model/vae/audio_vae.hpp" @@ -1492,6 +1493,11 @@ bool StableDiffusionGGML::build_denoiser() { break; } case H3_AV_FLOW_PRED: { + if (version == VERSION_MINIMAX_H3_IMAGE) { + LOG_INFO("running in MiniMax H3 Image FLOW mode"); + denoiser = std::make_shared(MiniMaxH3::pdd_sigmas); + break; + } 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; @@ -2615,7 +2621,7 @@ sd::Tensor StableDiffusionGGML::sample(const std::shared_ptr 0 && target_index < sigmas.size() ? sigmas[target_index] : NAN}; @@ -3054,7 +3060,7 @@ sd::Tensor StableDiffusionGGML::decode_first_stage(const sd::Tensordecode(n_threads, latents, tiling_params, decode_video, circular_x, circular_y); - if (version == VERSION_P_H3_IMAGE && !decoded.empty() && decoded.dim() == 5) { + if (sd_version_is_minimax_h3_image(version) && !decoded.empty() && decoded.dim() == 5) { decoded = decoded.squeeze(2); } if (!decoded.empty() || diff --git a/src/pipeline/image.cpp b/src/pipeline/image.cpp index 50a4fb7ac1..5e7a349583 100644 --- a/src/pipeline/image.cpp +++ b/src/pipeline/image.cpp @@ -172,6 +172,17 @@ namespace sd::pipeline { return packed; } + static sd::Tensor resize_h3_image_reference(const sd::Tensor& image, int64_t max_short_edge) { + const int64_t short_edge = std::min(image.shape()[0], image.shape()[1]); + const double scale = static_cast(std::clamp(short_edge, 1024, max_short_edge)) / short_edge; + const int64_t width = std::max(32, static_cast(std::nearbyint(image.shape()[0] * scale / 32)) * 32); + const 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 std::optional prepare_image_generation_latents(StableDiffusionGGML* sd, const sd_img_gen_params_t* sd_img_gen_params, GenerationRequest* request, @@ -354,7 +365,20 @@ namespace sd::pipeline { continue; } sd::Tensor ref_latent; - if (sd->version == VERSION_P_H3_IMAGE) { + if (sd->version == VERSION_MINIMAX_H3_IMAGE) { + ref_latent = sd->encode_first_stage(resize_h3_image_reference(ref_images[i], 2048)); + if (!ref_latent.empty()) { + const int64_t width = ref_latent.shape()[0]; + const int64_t height = ref_latent.shape()[1]; + const int64_t channels = ref_latent.shape()[3]; + auto condition_rng = std::make_shared(static_cast(request->seed)); + // Upstream augments patch rows, so the RNG follows token/channel/patch order. + auto noise = sd::randn({2, 2, channels, width / 2, height / 2}, condition_rng) + .permute({0, 3, 1, 4, 2}) + .reshape(ref_latent.shape()); + ref_latent = ref_latent * 0.999f + noise * (1.f - 0.999f); + } + } else 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)) { @@ -484,14 +508,18 @@ namespace sd::pipeline { } ImageGenerationLatents latents; - if (sd->version == VERSION_P_H3_IMAGE) { - init_latent = pack_ph3_image_latent(init_latent); + if (sd_version_is_minimax_h3_image(sd->version)) { + if (sd->version == VERSION_P_H3_IMAGE) { + init_latent = pack_ph3_image_latent(init_latent); + } else if (init_latent.dim() == 4) { + init_latent = init_latent.unsqueeze(2); + } 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); + image = sd->version == VERSION_P_H3_IMAGE ? resize_ph3_reference(image, 1000000) : resize_h3_image_reference(image, 1024); } } latents.init_latent = std::move(init_latent); @@ -548,7 +576,7 @@ 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) { + if (sd_version_is_minimax_h3_image(sd->version) && ref_image_params.pass_to_dit) { cond.c_reference_blocks = latents->minimax_reference_blocks; } @@ -596,7 +624,7 @@ 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) { + if (sd_version_is_minimax_h3_image(sd->version) && ref_image_params.pass_to_dit) { uncond.c_reference_blocks = latents->minimax_reference_blocks; } } diff --git a/src/pipeline/model_builders.cpp b/src/pipeline/model_builders.cpp index e78a0d666e..afc0ce8d7f 100644 --- a/src/pipeline/model_builders.cpp +++ b/src/pipeline/model_builders.cpp @@ -159,6 +159,14 @@ namespace sd::model_builders { version, weight_manager); } else if (find_model_traits(vae_version)->vae_kind == ModelVAEKind::MINIMAX_H3) { + if (version == VERSION_MINIMAX_H3_IMAGE) { + auto proj = tensor_storage_map.find("first_stage_model.decoder.proj_out.weight"); + auto embedder = tensor_storage_map.find("first_stage_model.decoder.x_embedder.weight"); + if (proj == tensor_storage_map.end() || proj->second.ne[1] != 768 || + embedder == tensor_storage_map.end() || embedder->second.ne[1] != 1024) { + throw std::runtime_error("MiniMax-H3-Image requires its 2D image VAE; pass --vae minimax_h3_image_vae_fp16.safetensors"); + } + } 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"); diff --git a/src/runtime/denoiser.hpp b/src/runtime/denoiser.hpp index 6fbcb43df3..701d8d9983 100644 --- a/src/runtime/denoiser.hpp +++ b/src/runtime/denoiser.hpp @@ -2,6 +2,7 @@ #define __SD_RUNTIME_DENOISER_HPP__ #include +#include #include #include #include @@ -1405,6 +1406,38 @@ struct MingImageFlowDenoiser : DiscreteFlowDenoiser { } }; +struct MiniMaxH3ImageDenoiser : public DiscreteFlowDenoiser { + std::array pdd_sigmas; + + explicit MiniMaxH3ImageDenoiser(const std::array& pdd_sigmas) + : DiscreteFlowDenoiser(1.f), pdd_sigmas(pdd_sigmas) {} + + float sigma_min() override { return 0.f; } + float sigma_max() override { return 1.f; } + + float t_to_sigma(float t) override { + const float position = (1.f - std::clamp(t / (TIMESTEPS - 1.f), 0.f, 1.f)) * 6.f; + const size_t index = std::min(static_cast(position), 5); + const auto& grid = pdd_sigmas; + return grid[index] + (position - index) * (grid[index + 1] - grid[index]); + } + + std::vector get_sigmas(uint32_t n, int image_seq_len, scheduler_t scheduler, SDVersion version, const char* extra_sample_args = nullptr) override { + if (scheduler != DISCRETE_SCHEDULER) { + return DiscreteFlowDenoiser::get_sigmas(n, image_seq_len, scheduler, version, extra_sample_args); + } + if (n == 0) { + return {}; + } + std::vector sigmas; + sigmas.reserve(n + 1); + for (uint32_t i = 0; i <= n; ++i) { + sigmas.push_back(t_to_sigma((1.f - static_cast(i) / n) * (TIMESTEPS - 1.f))); + } + return sigmas; + } +}; + struct H3AVFlowDenoiser : public DiscreteFlowDenoiser { int64_t video_channels; float audio_shift;