Commit 45a63ac for stable-diffusion.cpp
commit 45a63ac460b9403df7a739bbc8662cb7d867eace
Author: leejet <leejet714@gmail.com>
Date: Sun Nov 1 21:18:37 2026 +0800
refactor: move model-specific operators and fusions into backend extensions (#2253)
diff --git a/.gitmodules b/.gitmodules
index a26210d..9e1837e 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 d7bfda5..235556f 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 0000000..cf41592
--- /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 0000000..0ec6f32
--- /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 0000000..a664209
--- /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 <cstdint>\nnamespace sd_vulkan_shaders {\nstatic const uint32_t ${NAME}[] = {${SPIRV_WORDS}};\n}\n")
diff --git a/cmake/ggml.cmake b/cmake/ggml.cmake
index 39d77ea..55d0067 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 0000000..e0d8133
--- /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 0e7b1a9..62966ec 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 0000000..6787c44
--- /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 04d9c4f..e8036c6 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 46c8540..2c5d960 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 c57b71a..d49c49e 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 91c5f3b..479433b 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 568a0e5..29b3313 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 de1cea8..4e53175 100644
--- a/src/core/compute_workspace.cpp
+++ b/src/core/compute_workspace.cpp
@@ -6,9 +6,9 @@
#include <unordered_map>
#include <unordered_set>
+#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<size_t>(graph->n_nodes + graph->n_leafs) + 8;
+ const size_t capacity = static_cast<size_t>(ggml_graph_n_nodes(graph) + ggml_graph_n_leafs(graph)) + 8;
if (scheduler_ == nullptr || capacity > scheduler_capacity_) {
release();
scheduler_ = make_scheduler(capacity);
@@ -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<const uint8_t*>(&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<uintptr_t>(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<void*>(static_cast<uintptr_t>(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<void*>(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<const std::unordered_map<const ggml_tensor*, ggml_tensor*>*>(user_data);
+ const auto found = mapped.find(tensor);
+ return found == mapped.end() ? nullptr : found->second;
+ };
+ GGML_ASSERT(ggml_graph_cpy_mapped(graph, copy_graph, map_tensor, &copies));
Measurement result;
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 0000000..ec4be83
--- /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 <algorithm>
+#include <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+ static void transform(float* values, int group) {
+ const float scale = 1.f / std::sqrt(float(group));
+ for (int i = 0; i < group; ++i)
+ values[i] *= scale;
+ for (int stride = 1; stride < group; stride *= 4)
+ for (int base = 0; base < group; base += 4 * stride)
+ for (int j = 0; j < stride; ++j) {
+ const int i = base + j;
+ const float a = values[i], b = values[i + stride], c = values[i + 2 * stride], d = values[i + 3 * stride];
+ values[i] = a + b + c - d;
+ values[i + stride] = a + b - c + d;
+ values[i + 2 * stride] = a - b + c + d;
+ values[i + 3 * stride] = -a + b + c + d;
+ }
+ }
+ struct I8Execution {
+ SDI8Operation op;
+ int64_t rows;
+ int8_t* quant;
+ float* scales;
+ float* output;
+ };
+ static void quantize(ggml_tensor*, int ith, int nth, void* data) {
+ auto& e = *static_cast<I8Execution*>(data);
+ const int64_t k = e.op.input->ne[0];
+ const int group = e.op.group_size;
+ const auto* input = static_cast<const float*>(e.op.input->data);
+ for (int64_t row = ith; row < e.rows; row += nth) {
+ const float* src = input + row * k;
+ int8_t* dst = e.quant + row * k;
+ float amax = 0.f, values[256];
+ auto visit = [&](auto fn) { if (!group) { for (int64_t i = 0; i < k; ++i) fn(i, src[i]); }
+ else for (int64_t base = 0; base < k; base += group) { std::memcpy(values, src + base, group * sizeof(float)); transform(values, group); for (int j = 0; j < group; ++j) fn(base + j, values[j]); } };
+ visit([&](int64_t, float value) { amax = std::max(amax, std::fabs(value)); });
+ const float scale = amax / 127.f, inverse = scale == 0.f ? 0.f : 1.f / scale;
+ e.scales[row] = scale;
+ if (scale == 0.f)
+ std::memset(dst, 0, k);
+ else
+ visit([&](int64_t i, float value) { dst[i] = static_cast<int8_t>(std::max(-127L, std::min(127L, std::lrintf(value * inverse)))); });
+ }
+ }
+ static void multiply(ggml_tensor*, int ith, int nth, void* data) {
+ auto& e = *static_cast<I8Execution*>(data);
+ const auto& op = e.op;
+ const int64_t k = op.weight->ne[0], n = op.weight->ne[1];
+ const auto* w = static_cast<const int8_t*>(op.weight->data);
+ const auto* ws = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+ const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+ for (int64_t index = ith; index < e.rows * n; index += nth) {
+ const int64_t row = index / n, column = index % n;
+ uint32_t sum = 0;
+ for (int64_t i = 0; i < k; ++i)
+ sum += uint32_t(int32_t(w[column * k + i]) * int32_t(e.quant[row * k + i]));
+ int32_t signed_sum;
+ std::memcpy(&signed_sum, &sum, sizeof(sum));
+ float value = float(signed_sum) * e.scales[row];
+ if (ws)
+ value *= ws[column];
+ if (bias)
+ value += bias[column];
+ e.output[index] = value;
+ }
+ }
+ static bool overlaps(const ggml_tensor* dst, const ggml_tensor* src) {
+ if (!src)
+ return false;
+ const auto a = reinterpret_cast<uintptr_t>(dst->data), b = reinterpret_cast<uintptr_t>(src->data);
+ return a <= b ? b - a < ggml_nbytes(dst) : a - b < ggml_nbytes(src);
+ }
+ static size_t quant_workspace(const ggml_tensor* node, const SDI8Operation& op) {
+ return op.input->type == GGML_TYPE_F32 ? (size_t(ggml_nelements(op.input)) + 3) / 4 * 4 + size_t(ggml_nrows(node)) * 4 : 0;
+ }
+ static bool i8_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+ SDI8Operation op;
+ return sd_ggml_i8_describe(node, &op);
+ }
+ static bool convrot_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+ int group;
+ return sd_ggml_convrot_describe(node, &group);
+ }
+ static ggml_status i8_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+ SDI8Operation op;
+ if (!sd_ggml_i8_describe(nodes[0], &op))
+ return GGML_STATUS_FAILED;
+ *bytes = quant_workspace(nodes[0], op) + ggml_nbytes(nodes[0]);
+ return GGML_STATUS_SUCCESS;
+ }
+ static ggml_status convrot_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* bytes, void*) {
+ *bytes = ggml_nbytes(nodes[0]);
+ return GGML_STATUS_SUCCESS;
+ }
+ static void destroy(void*, void*) {}
+ static ggml_ext_cpu_parallel_fn parallel(ggml_backend_t backend) {
+ return reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)), "ggml_ext_cpu_parallel_v1"));
+ }
+ static ggml_status i8_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+ auto* dst = nodes[0];
+ I8Execution e = {};
+ if (!sd_ggml_i8_describe(dst, &e.op))
+ return GGML_STATUS_FAILED;
+ e.rows = ggml_nrows(dst);
+ auto run = parallel(context->backend);
+ if (!run)
+ return GGML_STATUS_FAILED;
+ char* scratch = static_cast<char*>(ggml_backend_buffer_get_base(context->workspace));
+ const size_t quant_size = quant_workspace(dst, e.op);
+ if (e.op.input->type == GGML_TYPE_F32) {
+ e.quant = reinterpret_cast<int8_t*>(scratch);
+ e.scales = reinterpret_cast<float*>(scratch + (size_t(ggml_nelements(e.op.input)) + 3) / 4 * 4);
+ auto status = run(context->backend, dst, quantize, &e);
+ if (status != GGML_STATUS_SUCCESS)
+ return status;
+ } else {
+ e.quant = static_cast<int8_t*>(e.op.input->data);
+ e.scales = e.op.packed ? reinterpret_cast<float*>(e.quant + e.op.input->ne[0] * ((e.rows + 3) / 4 * 4)) : static_cast<float*>(e.op.input_scale->data);
+ }
+ bool alias = false;
+ for (auto* input : dst->src)
+ alias |= overlaps(dst, input);
+ e.output = alias ? reinterpret_cast<float*>(scratch + quant_size) : static_cast<float*>(dst->data);
+ auto status = run(context->backend, dst, multiply, &e);
+ if (status == GGML_STATUS_SUCCESS && alias)
+ std::memcpy(dst->data, e.output, ggml_nbytes(dst));
+ return status;
+ }
+ static ggml_status convrot_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int, void*, void*) {
+ auto* dst = nodes[0];
+ I8Execution e = {};
+ if (!sd_ggml_convrot_describe(dst, &e.op.group_size))
+ return GGML_STATUS_FAILED;
+ auto run = parallel(context->backend);
+ if (!run)
+ return GGML_STATUS_FAILED;
+ e.op.input = dst->src[0];
+ e.rows = ggml_nrows(e.op.input);
+ bool alias = overlaps(dst, e.op.input);
+ e.quant = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(context->workspace) : dst->data);
+ const int64_t k = dst->ne[0], padded = (e.rows + 3) / 4 * 4;
+ e.scales = reinterpret_cast<float*>(e.quant + k * padded);
+ auto status = run(context->backend, dst, quantize, &e);
+ if (status == GGML_STATUS_SUCCESS) {
+ std::memset(e.quant + k * e.rows, 0, k * (padded - e.rows));
+ if (alias)
+ std::memcpy(dst->data, e.quant, ggml_nbytes(dst));
+ }
+ return status;
+ }
+ void register_i8(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc d = {};
+ d.abi_version = GGML_EXT_ABI_VERSION;
+ d.priority = 100;
+ d.destroy = destroy;
+ d.name = "sd.linear_i8.cpu.v1";
+ d.op = GGML_OP_CUSTOM;
+ d.custom_op = sd_ggml_i8_op();
+ d.supports = i8_supports;
+ d.prepare = i8_prepare;
+ d.execute = i8_execute;
+ GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+ d.name = "sd.quantize_i8_convrot.cpu.v1";
+ d.op = GGML_OP_CUSTOM;
+ d.custom_op = sd_ggml_convrot_op();
+ d.supports = convrot_supports;
+ d.prepare = convrot_prepare;
+ d.execute = convrot_execute;
+ GGML_ASSERT(ggml_ext_register_kernel(device, &d));
+ }
+
+}
+#endif
diff --git a/src/core/cpu/registration.cpp b/src/core/cpu/registration.cpp
new file mode 100644
index 0000000..a4f5bd3
--- /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 0000000..c711f6f
--- /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 0000000..ad1ac79
--- /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 <cmath>
+#include <cstring>
+
+namespace sd_cpu {
+
+ static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+ return sd_ggml_rope_describe(node);
+ }
+
+ struct RopeExecution {
+ const ggml_tensor* rope;
+ const ggml_tensor* norm;
+ const ggml_tensor* weight;
+ };
+
+ static void rope_compute(ggml_tensor* dst, int ith, int nth, void* user_data) {
+ const auto& execution = *static_cast<const RopeExecution*>(user_data);
+ const auto* rope = execution.rope;
+ const auto* x = execution.norm ? execution.norm->src[0] : rope->src[0];
+ const auto* pe = rope->src[1];
+ int32_t interleaved;
+ std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
+ float eps = 0.0f;
+ if (execution.norm) {
+ std::memcpy(&eps, execution.norm->op_params, sizeof(eps));
+ }
+ const int64_t half_dim = pe->ne[2];
+ const int64_t rows = x->ne[1] * x->ne[2] * x->ne[3];
+ for (int64_t row = ith; row < rows; row += nth) {
+ const int64_t token = row % x->ne[2];
+ const int64_t head = (row / x->ne[2]) % x->ne[1];
+ const int64_t batch = row / (x->ne[2] * x->ne[1]);
+ const char* input = static_cast<const char*>(x->data) + head * x->nb[1] + token * x->nb[2] + batch * x->nb[3];
+ const char* pos = static_cast<const char*>(pe->data) + token * pe->nb[3];
+ float* output = static_cast<float*>(dst->data) + row * x->ne[0];
+ float scale = 1.0f;
+ if (execution.norm) {
+ double sum = 0.0;
+ for (int64_t col = 0; col < x->ne[0]; ++col) {
+ const float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+ sum += static_cast<double>(v * v);
+ }
+ const float mean = static_cast<float>(sum / x->ne[0]);
+ scale = 1.0f / std::sqrt(mean + eps);
+ }
+ auto value = [&](int64_t col) {
+ float v = *reinterpret_cast<const float*>(input + col * x->nb[0]);
+ if (execution.norm) {
+ const auto* weight = static_cast<const char*>(execution.weight->data);
+ v = (v * scale) * *reinterpret_cast<const float*>(weight + col * execution.weight->nb[0]);
+ }
+ return v;
+ };
+ for (int64_t j = 0; j < half_dim; ++j) {
+ const int64_t i0 = interleaved ? 2 * j : j;
+ const int64_t i1 = interleaved ? 2 * j + 1 : j + half_dim;
+ const float a = value(i0);
+ const float b = value(i1);
+ const char* m = pos + j * pe->nb[2];
+ const float c = *reinterpret_cast<const float*>(m);
+ const float s = *reinterpret_cast<const float*>(m + pe->nb[1]);
+ if (interleaved) {
+ const float neg_s = *reinterpret_cast<const float*>(m + pe->nb[0]);
+ const float c1 = *reinterpret_cast<const float*>(m + pe->nb[0] + pe->nb[1]);
+ output[i0] = a * c + b * neg_s;
+ output[i1] = a * s + b * c1;
+ } else {
+ output[i0] = a * c - b * s;
+ output[i1] = a * s + b * c;
+ }
+ }
+ for (int64_t col = 2 * half_dim; col < x->ne[0]; ++col) {
+ output[col] = value(col);
+ }
+ }
+ }
+
+ static int rope_fusion_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+ auto* view = const_cast<ggml_cgraph*>(graph);
+ if (first + 2 >= ggml_graph_n_nodes(view)) {
+ return 0;
+ }
+ const auto* norm = ggml_graph_node(view, first);
+ const auto* mul = ggml_graph_node(view, first + 1);
+ const auto* rope = ggml_graph_node(view, first + 2);
+ if (norm->op != GGML_OP_RMS_NORM || norm->type != GGML_TYPE_F32 || !norm->src[0] ||
+ norm->src[0]->type != GGML_TYPE_F32 || norm->src[0]->nb[0] != sizeof(float) ||
+ mul->op != GGML_OP_MUL || mul->type != GGML_TYPE_F32 || rope->src[0] != mul ||
+ ggml_ext_op_get(rope) != sd_ggml_rope_op() ||
+ !sd_ggml_rope_describe(rope) || !ggml_are_same_shape(norm, mul)) {
+ return 0;
+ }
+ const auto* weight = mul->src[0] == norm ? mul->src[1] : (mul->src[1] == norm ? mul->src[0] : nullptr);
+ if (!weight || weight->type != GGML_TYPE_F32 || weight->ne[0] != norm->ne[0] ||
+ weight->ne[1] != 1 || weight->ne[2] != 1 || weight->ne[3] != 1) {
+ return 0;
+ }
+ return 3;
+ }
+
+ static ggml_status rope_execute(const ggml_ext_execution_context* context, ggml_tensor* const* nodes, int count, void*, void*) {
+ auto device = ggml_backend_get_device(context->backend);
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto run = reinterpret_cast<ggml_ext_cpu_parallel_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cpu_parallel_v1"));
+ GGML_ASSERT(run);
+ RopeExecution execution = {nodes[count - 1], nullptr, nullptr};
+ if (count == 3) {
+ execution.norm = nodes[0];
+ execution.weight = nodes[1]->src[0] == nodes[0] ? nodes[1]->src[1] : nodes[1]->src[0];
+ }
+ return run(context->backend, nodes[count - 1], rope_compute, &execution);
+ }
+
+ void register_rope(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.rope_apply.cpu.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_rope_op();
+ desc.priority = 100;
+ desc.supports = rope_supports;
+ desc.execute = rope_execute;
+ GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+
+ desc.name = "sd.rms_norm_mul_rope.cpu.v1";
+ desc.op = GGML_OP_RMS_NORM;
+ desc.custom_op = nullptr;
+ desc.priority = 200;
+ desc.supports = nullptr;
+ desc.match = rope_fusion_match;
+ GGML_ASSERT(ggml_ext_register_kernel(device, &desc));
+ }
+
+}
+#endif
diff --git a/src/core/cuda/CMakeLists.txt b/src/core/cuda/CMakeLists.txt
new file mode 100644
index 0000000..7bca6aa
--- /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 0000000..2cd660b
--- /dev/null
+++ b/src/core/cuda/affine.cu
@@ -0,0 +1,183 @@
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include "ggml-ext.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+ namespace sd_affine_cuda {
+ static bool disabled() {
+ static const bool value = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+ return value;
+ }
+ static bool aligned(const ggml_tensor* tensor, size_t alignment) {
+ size_t offset = 0;
+ while (!tensor->data && tensor->view_src) {
+ offset += tensor->view_offs;
+ tensor = tensor->view_src;
+ }
+ return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+ }
+ static float param(const ggml_tensor* tensor, int index) {
+ float value;
+ std::memcpy(&value, tensor->op_params + index, sizeof(value));
+ return value;
+ }
+ static bool contiguous_f32(const ggml_tensor* tensor) {
+ return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor) && aligned(tensor, 16);
+ }
+ static bool bounded(const ggml_tensor* tensor) {
+ return ggml_nelements(tensor) > 0 && ggml_nelements(tensor) <= int64_t(INT_MAX) * 1024;
+ }
+ static int match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+ auto* g = const_cast<ggml_cgraph*>(graph);
+ if (disabled() || first + 1 >= ggml_graph_n_nodes(g))
+ return 0;
+ const auto* node = ggml_graph_node(g, first);
+ const auto* next = ggml_graph_node(g, first + 1);
+ if (node->op == GGML_OP_MUL) {
+ if (next->op != GGML_OP_ADD || (next->src[0] != node && next->src[1] != node))
+ return 0;
+ const auto* x = node->src[0];
+ const auto* gate = node->src[1];
+ const auto* residual = next->src[0] == node ? next->src[1] : next->src[0];
+ return residual != node && contiguous_f32(x) && contiguous_f32(gate) && contiguous_f32(residual) && contiguous_f32(next) &&
+ node->type == GGML_TYPE_F32 && ggml_are_same_shape(node, next) && ggml_are_same_shape(x, next) &&
+ ggml_are_same_shape(residual, next) && gate->ne[0] == x->ne[0] && ggml_nelements(gate) == x->ne[0] &&
+ x->ne[0] % 4 == 0 && bounded(next)
+ ? 2
+ : 0;
+ }
+ if (node->op == GGML_OP_CONCAT) {
+ const bool permuted = next->op == GGML_OP_PERMUTE;
+ const int count = permuted ? 3 : 2;
+ if (first + count > ggml_graph_n_nodes(g) || node->type != GGML_TYPE_F32 || node->ne[0] % 4 ||
+ node->op_params[0] != (permuted ? 2 : 1) || !ggml_is_contiguous(node) ||
+ !contiguous_f32(node->src[0]) || !contiguous_f32(node->src[1]))
+ return 0;
+ if (permuted && (next->src[0] != node || next->ne[0] != node->ne[0] || next->ne[1] != node->ne[2] ||
+ next->ne[2] != node->ne[1] || next->ne[3] != node->ne[3] || next->nb[0] != node->nb[0] ||
+ next->nb[1] != node->nb[2] || next->nb[2] != node->nb[1] || next->nb[3] != node->nb[3]))
+ return 0;
+ const auto* dst = ggml_graph_node(g, first + count - 1);
+ return dst->op == GGML_OP_CPY && dst->src[0] == (permuted ? next : node) && dst->type == GGML_TYPE_F16 &&
+ ggml_is_contiguous(dst) && ggml_are_same_shape(dst, dst->src[0]) && aligned(dst, 4) && bounded(dst)
+ ? count
+ : 0;
+ }
+ if (node->op == GGML_OP_CONT && first + 3 < ggml_graph_n_nodes(g)) {
+ const auto* src = node->src[0];
+ const auto* permute = ggml_graph_node(g, first + 2);
+ const auto* dst = ggml_graph_node(g, first + 3);
+ return ggml_are_same_shape(src, node) && src->type == GGML_TYPE_F32 && src->nb[0] == 4 && src->ne[0] % 4 == 0 &&
+ aligned(src, 16) && (src->nb[1] | src->nb[2] | src->nb[3]) % 16 == 0 &&
+ next->op == GGML_OP_SCALE && next->src[0] == node && next->type == GGML_TYPE_F32 && param(next, 1) == 0.f &&
+ permute->op == GGML_OP_PERMUTE && permute->src[0] == next && permute->op_params[0] == 0 &&
+ permute->op_params[1] == 2 && permute->op_params[2] == 1 && permute->op_params[3] == 3 &&
+ dst->op == GGML_OP_CPY && dst->src[0] == permute && dst->type == GGML_TYPE_F16 &&
+ ggml_are_same_shape(dst, permute) && ggml_is_contiguous(dst) && aligned(dst, 4) && bounded(dst)
+ ? 4
+ : 0;
+ }
+ return 0;
+ }
+ static __global__ void gated_residual(const float4* x, const float4* gate, const float4* residual, float4* dst, int64_t n4, int64_t count4) {
+ const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (i >= count4)
+ return;
+ const float4 a = x[i], g = gate[i % n4], r = residual[i];
+ // The graph performs two FP32 operations; contracting them into FMA changes rounding.
+ dst[i] = make_float4(__fadd_rn(__fmul_rn(a.x, g.x), r.x), __fadd_rn(__fmul_rn(a.y, g.y), r.y),
+ __fadd_rn(__fmul_rn(a.z, g.z), r.z), __fadd_rn(__fmul_rn(a.w, g.w), r.w));
+ }
+ template <bool Permuted>
+ static __global__ void concat_cast(const float4* prefix, const float4* current, half2* dst, int64_t columns4, int64_t prefix_tokens, int64_t tokens, int64_t heads, int64_t count4) {
+ const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (i >= count4)
+ return;
+ const int64_t col = i % columns4, row = i / columns4;
+ const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
+ const bool from_prefix = token < prefix_tokens;
+ const int64_t src_tokens = from_prefix ? prefix_tokens : tokens - prefix_tokens;
+ const int64_t src_token = from_prefix ? token : token - prefix_tokens;
+ const int64_t src_row = Permuted ? (batch * src_tokens + src_token) * heads + head
+ : (batch * heads + head) * src_tokens + src_token;
+ const float4 value = (from_prefix ? prefix : current)[src_row * columns4 + col];
+ dst[2 * i] = __floats2half2_rn(value.x, value.y);
+ dst[2 * i + 1] = __floats2half2_rn(value.z, value.w);
+ }
+ static __global__ void scale_permute(const float* src, half2* dst, int64_t dim, int64_t tokens, int64_t heads, int64_t count, int64_t nb1, int64_t nb2, int64_t nb3, float scale) {
+ const int64_t index = (int64_t(blockIdx.x) * blockDim.x + threadIdx.x) * 4;
+ if (index >= count)
+ return;
+ const int64_t col = index % dim, row = index / dim;
+ const int64_t token = row % tokens, head = row / tokens % heads, batch = row / (tokens * heads);
+ const float4 value = *reinterpret_cast<const float4*>(src + col + head * nb1 + token * nb2 + batch * nb3);
+ dst[index / 2] = __floats2half2_rn(__fmul_rn(value.x, scale), __fmul_rn(value.y, scale));
+ dst[index / 2 + 1] = __floats2half2_rn(__fmul_rn(value.z, scale), __fmul_rn(value.w, scale));
+ }
+ static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+ }
+ static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+ *workspace = ggml_nbytes(nodes[count - 1]);
+ return GGML_STATUS_SUCCESS;
+ }
+ static void destroy(void*, void*) {}
+ static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<CUDAStream>(context.stream);
+ const auto* node = nodes[0];
+ const auto* a = node->src[0];
+ const auto* b = node->op == GGML_OP_CONT ? nullptr : node->src[1];
+ auto* dst = nodes[count - 1];
+ const auto* residual = node->op == GGML_OP_MUL ? (dst->src[0] == node ? dst->src[1] : dst->src[0]) : nullptr;
+ const bool gate = residual != nullptr;
+ const bool alias = (overlaps(dst, a) && !(gate && dst->data == a->data)) || (b && overlaps(dst, b)) ||
+ (residual && overlaps(dst, residual) && dst->data != residual->data);
+ void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+ const int64_t count4 = ggml_nelements(dst) / 4;
+ const unsigned blocks = unsigned((count4 + 255) / 256);
+ if (gate) {
+ gated_residual<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
+ static_cast<const float4*>(residual->data), static_cast<float4*>(output), a->ne[0] / 4, count4);
+ } else if (node->op == GGML_OP_CONCAT) {
+ const bool permuted = count == 3;
+ auto kernel = permuted ? concat_cast<true> : concat_cast<false>;
+ kernel<<<blocks, 256, 0, stream>>>(static_cast<const float4*>(a->data), static_cast<const float4*>(b->data),
+ static_cast<half2*>(output), node->ne[0] / 4, a->ne[permuted ? 2 : 1], node->ne[permuted ? 2 : 1],
+ node->ne[permuted ? 1 : 2], count4);
+ } else {
+ scale_permute<<<blocks, 256, 0, stream>>>(static_cast<const float*>(a->data), static_cast<half2*>(output),
+ a->ne[0], a->ne[2], a->ne[1], ggml_nelements(dst), a->nb[1] / 4, a->nb[2] / 4, a->nb[3] / 4, param(nodes[1], 0));
+ }
+ auto error = cuda_last_error();
+ if (error == cuda_success && alias)
+ error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+ return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+ }
+
+ void register_affine(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.priority = 200;
+ desc.match = sd_affine_cuda::match;
+ desc.prepare = sd_affine_cuda::prepare;
+ desc.destroy = sd_affine_cuda::destroy;
+ desc.execute = sd_affine_cuda::execute;
+ desc.name = "sd.gated_residual." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_MUL;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.concat_cast." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_CONCAT;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.cont_scale_permute_cast." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_CONT;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ }
+
+}
diff --git a/src/core/cuda/convrot.cu b/src/core/cuda/convrot.cu
new file mode 100644
index 0000000..d5d159b
--- /dev/null
+++ b/src/core/cuda/convrot.cu
@@ -0,0 +1,296 @@
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <string>
+#include "convrot_kernel.h"
+#include "core/ggml_extensions.h"
+#include "runtime.h"
+
+namespace SD_CUDA_NAMESPACE {
+ namespace sd_convrot {
+
+ static bool fusion_disabled() {
+ static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+ return disabled;
+ }
+ static float param_f32(const ggml_tensor* tensor, int index) {
+ float value;
+ std::memcpy(&value, tensor->op_params + index, sizeof(value));
+ return value;
+ }
+ static bool f32(const ggml_tensor* tensor) {
+ return tensor && tensor->type == GGML_TYPE_F32 && ggml_is_contiguous(tensor);
+ }
+ static bool vector(const ggml_tensor* tensor, int64_t k) {
+ return f32(tensor) && tensor->ne[0] == k && ggml_nelements(tensor) == k;
+ }
+ enum Kind { SCALE,
+ NORM,
+ 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<ggml_cgraph*>(graph);
+ const int available = ggml_graph_n_nodes(g) - first;
+ const auto* node = ggml_graph_node(g, first);
+ const int count = node->op == GGML_OP_RMS_NORM ? 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<Plan*>(state);
+ }
+ static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ if (!b)
+ return false;
+ const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+ }
+ template <bool Swiglu = false, bool Normalize = false, bool Scaled = false, bool Sigmoid = false, bool RMS = false>
+ 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<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS>,
+ int(p.device.shared_memory - extra));
+ if (status != cuda_success)
+ return status;
+ p.configured = true;
+ }
+ const int threads = Normalize ? 1024 : RMS ? (op.k <= 1024 ? 128 : op.k <= 6144 ? 256
+ : op.k <= 8192 ? 512
+ : 1024)
+ : Swiglu && !Sigmoid && op.k >= 16384 ? 512
+ : 256;
+ const int64_t padded = (op.rows + 3) / 4 * 4;
+ auto data = [](const ggml_tensor* t) { return t ? static_cast<const float*>(t->data) : nullptr; };
+ quantize<Swiglu, false, Normalize, Scaled, warp, float, Sigmoid, RMS><<<unsigned(op.rows), threads, size_t(op.k) * sizeof(float), stream>>>(
+ data(op.input) + op.gate_offset, output, reinterpret_cast<float*>(output + op.k * padded), op.k, op.rows, op.k,
+ op.up ? data(op.up) + op.up_offset : nullptr, op.stride, op.up_stride, nullptr, nullptr, data(op.modulation), op.eps, op.scale, op.bias, data(op.weight), data(op.shift));
+ return cuda_last_error();
+ }
+ static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+ auto& plan = *static_cast<Plan*>(state);
+ const auto& op = plan.op;
+ auto* dst = nodes[count - 1];
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context) || context.device != plan.device.device)
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<CUDAStream>(context.stream);
+ bool alias = false;
+ for (const auto* input : {op.input, op.up, op.modulation, op.weight, op.shift})
+ alias |= overlaps(dst, input);
+ auto* output = static_cast<int8_t*>(alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data);
+ const int64_t padded = (op.rows + 3) / 4 * 4;
+ if (padded > op.rows && cuda_memset_async(output + op.k * op.rows, 0, size_t(op.k) * (padded - op.rows), stream) != cuda_success)
+ return GGML_STATUS_FAILED;
+ CUDAError error;
+ switch (op.kind) {
+ case SCALE:
+ error = launch<false, false, true>(plan, op, output, stream);
+ break;
+ case NORM:
+ error = launch<false, true>(plan, op, output, stream);
+ break;
+ case RMS_MOD:
+ error = launch<false, false, false, false, true>(plan, op, output, stream);
+ break;
+ case SWIGLU:
+ case SILU:
+ error = launch<true>(plan, op, output, stream);
+ break;
+ case SIGMOID:
+ error = launch<true, false, false, true>(plan, op, output, stream);
+ break;
+ default:
+ return GGML_STATUS_FAILED;
+ }
+ if (error != cuda_success || (alias && cuda_copy_async(dst->data, output, size_t(op.k) * padded + op.rows * sizeof(float), stream) != cuda_success))
+ return GGML_STATUS_FAILED;
+ return GGML_STATUS_SUCCESS;
+ }
+ }
+
+ void register_convrot(ggml_backend_dev_t device) {
+ for (auto op : {GGML_OP_SCALE, GGML_OP_NORM, GGML_OP_RMS_NORM, GGML_OP_GLU, GGML_OP_UNARY}) {
+ ggml_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 0000000..748c08e
--- /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 <cstdint>
+#include "runtime.h"
+namespace SD_CUDA_NAMESPACE {
+ namespace sd_convrot {
+ static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+ for (int offset = 16; offset > 0; offset >>= 1)
+ x += cuda_shuffle_xor(x, offset, 32);
+ return x;
+ }
+ static __device__ __forceinline__ float2 warp_reduce_sum(float2 x) {
+#pragma unroll
+ for (int offset = 16; offset > 0; offset >>= 1) {
+ x.x += cuda_shuffle_xor(x.x, offset, 32);
+ x.y += cuda_shuffle_xor(x.y, offset, 32);
+ }
+ return x;
+ }
+ template <int Width>
+ static __device__ __forceinline__ float warp_reduce_max(float x) {
+#pragma unroll
+ for (int offset = Width / 2; offset > 0; offset >>= 1)
+ x = fmaxf(x, cuda_shuffle_xor(x, offset, Width));
+ return x;
+ }
+ static __device__ float2 block_sum(float2 x, float2* shared) {
+ x = warp_reduce_sum(x);
+ if (threadIdx.x % 32 == 0)
+ shared[threadIdx.x / 32] = x;
+ __syncthreads();
+ return warp_reduce_sum(shared[threadIdx.x % 32]);
+ }
+ static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+ a = index == 3 ? -a : a;
+ b = index == 2 ? -b : b;
+ c = index == 1 ? -c : c;
+ d = index == 0 ? -d : d;
+ return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+ }
+
+ template <bool swiglu, bool dequantize = false, bool normalize = false, bool scaled = false, bool warp_rotation = false, typename T = float, bool sigmoid_gate = false, bool rms_modulation = false>
+ __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 0000000..cdafc99
--- /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 <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cutlass/numeric_conversion.h>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cute/tensor.hpp>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+#include "core/cuda/i8_projection.h"
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+#include "ggml-ext.h"
+
+static __device__ __forceinline__ float sd_warp_max(float v) {
+ for (int mask = 16; mask > 0; mask >>= 1)
+ v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, mask));
+ return v;
+}
+namespace sd_fattn_cute {
+ using namespace cute;
+ template <bool A_in_regs = false, bool B_in_regs = false, typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename Tensor4, typename TiledMma, typename TiledCopyA, typename TiledCopyB, typename ThrCopyA, typename ThrCopyB>
+ __forceinline__ __device__ void gemm(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsA, Tensor4 const& tCsB, TiledMma tiled_mma, TiledCopyA smem_tiled_copy_A, TiledCopyB smem_tiled_copy_B, ThrCopyA smem_thr_copy_A, ThrCopyB smem_thr_copy_B) {
+ CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+ CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+ CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+ Tensor tCrA_copy_view = smem_thr_copy_A.retile_D(tCrA);
+ CUTE_STATIC_ASSERT_V(size<1>(tCsA) == size<1>(tCrA_copy_view));
+ Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+ CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+ if (!A_in_regs) {
+ cute::copy(smem_tiled_copy_A, tCsA(_, _, _0{}), tCrA_copy_view(_, _, _0{}));
+ }
+ if (!B_in_regs) {
+ cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+ }
+#pragma unroll
+ for (int i = 0; i < size<2>(tCrA); ++i) {
+ if (i < size<2>(tCrA) - 1) {
+ if (!A_in_regs) {
+ cute::copy(smem_tiled_copy_A, tCsA(_, _, i + 1), tCrA_copy_view(_, _, i + 1));
+ }
+ if (!B_in_regs) {
+ cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+ }
+ }
+ cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+ }
+ }
+
+ template <typename Tensor0, typename Tensor1, typename Tensor2, typename Tensor3, typename TiledMma, typename TiledCopy, typename ThrCopy>
+ __forceinline__ __device__ void gemm_rs(Tensor0& acc, Tensor1& tCrA, Tensor2& tCrB, Tensor3 const& tCsB, TiledMma tiled_mma, TiledCopy smem_tiled_copy_B, ThrCopy smem_thr_copy_B) {
+ CUTE_STATIC_ASSERT_V(size<1>(tCrA) == size<1>(acc));
+ CUTE_STATIC_ASSERT_V(size<1>(tCrB) == size<2>(acc));
+ CUTE_STATIC_ASSERT_V(size<2>(tCrA) == size<2>(tCrB));
+ Tensor tCrB_copy_view = smem_thr_copy_B.retile_D(tCrB);
+ CUTE_STATIC_ASSERT_V(size<1>(tCsB) == size<1>(tCrB_copy_view));
+ cute::copy(smem_tiled_copy_B, tCsB(_, _, _0{}), tCrB_copy_view(_, _, _0{}));
+#pragma unroll
+ for (int i = 0; i < size<2>(tCrA); ++i) {
+ if (i < size<2>(tCrA) - 1) {
+ cute::copy(smem_tiled_copy_B, tCsB(_, _, i + 1), tCrB_copy_view(_, _, i + 1));
+ }
+ cute::gemm(tiled_mma, tCrA(_, _, i), tCrB(_, _, i), acc);
+ }
+ }
+
+ template <typename Layout>
+ __forceinline__ __device__ auto convert_layout_acc_rowcol(Layout acc_layout) {
+ static_assert(decltype(size<0>(acc_layout))::value == 4);
+ static_assert(decltype(rank(acc_layout))::value == 3);
+ auto l = logical_divide(acc_layout, Shape<_2>{});
+ return make_layout(make_layout(get<0, 1>(l), get<1>(l)), make_layout(get<0, 0>(l), get<2>(l)));
+ };
+
+ template <typename MMA_traits, typename Layout>
+ __forceinline__ __device__ auto convert_layout_acc_Aregs(Layout acc_layout) {
+ using X = Underscore;
+ static_assert(decltype(size<0>(acc_layout))::value == 4);
+ static_assert(decltype(rank(acc_layout))::value == 3);
+ constexpr int mma_shape_K = get<2>(typename MMA_traits::Shape_MNK{});
+ static_assert(mma_shape_K == 8 || mma_shape_K == 16);
+ if constexpr (mma_shape_K == 8) {
+ return acc_layout;
+ } else {
+ auto l = logical_divide(acc_layout, Shape<X, X, _2>{});
+ return make_layout(make_layout(get<0>(l), get<2, 0>(l)), get<1>(l), get<2, 1>(l));
+ }
+ };
+
+ template <typename TensorT>
+ __device__ __forceinline__ auto to_half(const TensorT& src) {
+ constexpr int count = decltype(size(src))::value;
+ auto dst = make_tensor<cutlass::half_t>(src.layout());
+ cutlass::NumericArrayConverter<cutlass::half_t, float, count> convert;
+ *reinterpret_cast<cutlass::Array<cutlass::half_t, count>*>(dst.data()) =
+ convert(*reinterpret_cast<const cutlass::Array<float, count>*>(src.data()));
+ return dst;
+ }
+
+ template <bool Full, bool Clear, bool FullCols = true, typename Copy, typename Src, typename Dst, typename Id>
+ __device__ __forceinline__ void copy_rows(Copy copier, const Src& src, Dst& dst, const Id& id, int rows, int cols = 0) {
+#pragma unroll
+ for (int m = 0; m < size<1>(src); ++m) {
+#pragma unroll
+ for (int k = 0; k < size<2>(src); ++k) {
+ if ((Full || get<0>(id(0, m, k)) < rows) && (FullCols || get<1>(id(0, m, k)) < cols)) {
+ cute::copy(copier, src(_, m, k), dst(_, m, k));
+ } else if constexpr (Clear) {
+ cute::clear(dst(_, m, k));
+ }
+ }
+ }
+ }
+
+ template <int Rows>
+ struct Softmax {
+ decltype(make_tensor<float>(Shape<Int<Rows>>{})) maximum, sum;
+ template <bool First, typename Scores, typename Output>
+ __device__ __forceinline__ void update(Scores& acc_s, Output& acc_o) {
+ auto scores = make_tensor(acc_s.data(), convert_layout_acc_rowcol(acc_s.layout()));
+ auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+ for (int m = 0; m < Rows; ++m) {
+ const float old_max = First ? -INFINITY : maximum(m);
+ float next_max = First ? scores(m, 0) : fmaxf(old_max, scores(m, 0));
+#pragma unroll
+ for (int n = 1; n < size<1>(scores); ++n)
+ next_max = fmaxf(next_max, scores(m, n));
+ next_max = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 2));
+ next_max = fmaxf(next_max, __shfl_xor_sync(0xffffffff, next_max, 1));
+ maximum(m) = next_max;
+ float row_sum = 0.f;
+ if constexpr (!First) {
+ row_sum = sum(m);
+ if (old_max != next_max) {
+ const float factor = exp2f((old_max - next_max) * 1.4426950408889634f);
+ row_sum *= factor;
+#pragma unroll
+ for (int n = 0; n < size<1>(output); ++n)
+ output(m, n) *= factor;
+ }
+ }
+#pragma unroll
+ for (int n = 0; n < size<1>(scores); ++n) {
+ const float 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 <typename Output>
+ __device__ __forceinline__ void normalize(Output& acc_o) {
+ auto output = make_tensor(acc_o.data(), convert_layout_acc_rowcol(acc_o.layout()));
+#pragma unroll
+ for (int m = 0; m < Rows; ++m) {
+ float total = sum(m);
+ total += __shfl_xor_sync(0xffffffff, total, 2);
+ total += __shfl_xor_sync(0xffffffff, total, 1);
+ const float inverse = 1.f / total;
+#pragma unroll
+ for (int n = 0; n < size<1>(output); ++n)
+ output(m, n) *= inverse;
+ }
+ }
+ };
+
+ template <int BM, int BN, int D = 128, int Warps = 4>
+ struct Traits {
+ using Element = cutlass::half_t;
+ static constexpr int Threads = Warps * 32;
+ using Mma = TiledMMA<MMA_Atom<SM80_16x8x16_F32F16F16F32_TN>,
+ Layout<Shape<Int<Warps>, _1, _1>>,
+ Tile<Int<16 * Warps>, _16, _16>>;
+ static constexpr int SmemWidth = D % 64 == 0 ? 64 : D % 32 == 0 ? 32
+ : 16;
+ static constexpr int SwizzleBits = SmemWidth == 64 ? 3 : SmemWidth == 32 ? 2
+ : 1;
+ using SmemAtom = decltype(composition(Swizzle<SwizzleBits, 3, 3>{},
+ Layout<Shape<_8, Int<SmemWidth>>, Stride<Int<SmemWidth>, _1>>{}));
+ using QLayout = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BM>, Int<D>>{}));
+ using KVLayout = decltype(tile_to_shape(SmemAtom{}, Shape<Int<BN>, Int<D>>{}));
+ using VtLayout = decltype(composition(KVLayout{}, make_layout(Shape<Int<D>, Int<BN>>{}, GenRowMajor{})));
+ using CopyAtom = Copy_Atom<SM75_U32x4_LDSM_N, Element>;
+ using CopyAtomT = Copy_Atom<SM75_U16x8_LDSM_T, Element>;
+ using GlobalLayout = Layout<Shape<Int<Threads / (SmemWidth / 8)>, Int<SmemWidth / 8>>, Stride<Int<SmemWidth / 8>, _1>>;
+ using GlobalCopy = decltype(make_tiled_copy(Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, Element>{},
+ GlobalLayout{},
+ Layout<Shape<_1, _8>>{}));
+ using OutputCopy = decltype(make_tiled_copy(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{},
+ GlobalLayout{},
+ Layout<Shape<_1, _8>>{}));
+ static constexpr int SharedBytes = (BM + 2 * BN) * D * sizeof(Element);
+ };
+
+ struct Params {
+ const cutlass::half_t *q, *k, *v;
+ cutlass::half_t* o;
+ int nq, nk, heads, kv_heads, dim;
+ };
+
+ template <int BM, int BN, int D = 128, bool FullDim = true, int Warps = 4>
+ __global__ __launch_bounds__(Warps * 32) void kernel(Params p) {
+ using T = Traits<BM, BN, D, Warps>;
+ using Element = typename T::Element;
+ extern __shared__ char storage[];
+ const int tid = threadIdx.x, m_block = blockIdx.x, batch = blockIdx.y, head = blockIdx.z;
+ const int64_t q_offset = (int64_t(batch) * p.heads + head) * p.nq * p.dim;
+ const int64_t kv_offset = (int64_t(batch) * p.kv_heads + head / (p.heads / p.kv_heads)) * p.nk * p.dim;
+ auto mQ = make_tensor(make_gmem_ptr(p.q + q_offset), make_shape(p.nq, p.dim), make_stride(p.dim, _1{}));
+ auto mK = make_tensor(make_gmem_ptr(p.k + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
+ auto mV = make_tensor(make_gmem_ptr(p.v + kv_offset), make_shape(p.nk, p.dim), make_stride(p.dim, _1{}));
+ auto gQ = local_tile(mQ, Shape<Int<BM>, Int<D>>{}, make_coord(m_block, 0));
+ auto gK = local_tile(mK, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
+ auto gV = local_tile(mV, Shape<Int<BN>, Int<D>>{}, make_coord(_, 0));
+ auto sQ = make_tensor(make_smem_ptr(reinterpret_cast<Element*>(storage)), typename T::QLayout{});
+ auto sK = make_tensor(sQ.data() + size(sQ), typename T::KVLayout{});
+ auto sV = make_tensor(sK.data() + size(sK), typename T::KVLayout{});
+ auto sVt = make_tensor(sV.data(), typename T::VtLayout{});
+ auto sVplain = make_tensor(sV.data().get(), get_nonswizzle_portion(typename T::VtLayout{}));
+ typename T::GlobalCopy global_copy;
+ auto thread_copy = global_copy.get_thread_slice(tid);
+ auto tQgQ = thread_copy.partition_S(gQ);
+ auto tQsQ = thread_copy.partition_D(sQ);
+ auto tKgK = thread_copy.partition_S(gK);
+ auto tKsK = thread_copy.partition_D(sK);
+ auto tVgV = thread_copy.partition_S(gV);
+ auto tVsV = thread_copy.partition_D(sV);
+ auto tQcQ = thread_copy.partition_S(make_identity_tensor(Shape<Int<BM>, Int<D>>{}));
+ auto tKVcKV = thread_copy.partition_S(make_identity_tensor(Shape<Int<BN>, Int<D>>{}));
+ typename T::Mma mma;
+ auto thread_mma = mma.get_thread_slice(tid);
+ auto rQ = thread_mma.partition_fragment_A(sQ);
+ auto rK = thread_mma.partition_fragment_B(sK);
+ auto rV = thread_mma.partition_fragment_B(sVplain);
+ auto acc_o = partition_fragment_C(mma, Shape<Int<BM>, Int<D>>{});
+ auto copyQ = make_tiled_copy_A(typename T::CopyAtom{}, mma);
+ auto copyK = make_tiled_copy_B(typename T::CopyAtom{}, mma);
+ auto copyV = make_tiled_copy_B(typename T::CopyAtomT{}, mma);
+ auto threadQ = copyQ.get_thread_slice(tid);
+ auto threadK = copyK.get_thread_slice(tid);
+ auto threadV = copyV.get_thread_slice(tid);
+ auto tSsQ = threadQ.partition_S(sQ);
+ auto tSsK = threadK.partition_S(sK);
+ auto tOsV = threadV.partition_S(sVt);
+ auto score_id = thread_mma.partition_C(make_identity_tensor(Shape<Int<BM>, Int<BN>>{}));
+ // Process the partial K/V tile first; subsequent tiles need no score masking.
+ int block = (p.nk - 1) / BN;
+ copy_rows<false, true, FullDim>(global_copy, tQgQ, tQsQ, tQcQ, p.nq - m_block * BM, p.dim);
+ auto last_k = tKgK(_, _, _, block);
+ copy_rows<false, true, FullDim>(global_copy, last_k, tKsK, tKVcKV, p.nk - block * BN, p.dim);
+ cute::cp_async_fence();
+ clear(acc_o);
+ Softmax<2 * size<1>(acc_o)> softmax;
+ auto step = [&](auto first_flag, int n) {
+ constexpr bool First = decltype(first_flag)::value;
+ auto acc_s = partition_fragment_C(mma, Shape<Int<BM>, Int<BN>>{});
+ clear(acc_s);
+ cute::cp_async_wait<0>();
+ __syncthreads();
+ auto v_tile = tVgV(_, _, _, n);
+ copy_rows<!First, First, FullDim>(global_copy, v_tile, tVsV, tKVcKV, p.nk - n * BN, p.dim);
+ cute::cp_async_fence();
+ gemm(acc_s, rQ, rK, tSsQ, tSsK, mma, copyQ, copyK, threadQ, threadK);
+ if constexpr (First) {
+#pragma unroll
+ for (int i = 0; i < size(acc_s); ++i) {
+ if (get<1>(score_id(i)) >= p.nk - n * BN)
+ acc_s(i) = -INFINITY;
+ }
+ }
+ cute::cp_async_wait<0>();
+ __syncthreads();
+ // Do not commit an empty async group: the next iteration waits for this K tile.
+ if (n > 0) {
+ auto next_k = tKgK(_, _, _, n - 1);
+ copy_rows<true, false, FullDim>(global_copy, next_k, tKsK, tKVcKV, 0, p.dim);
+ cute::cp_async_fence();
+ }
+ softmax.template update<First>(acc_s, acc_o);
+ auto rP = to_half(acc_s);
+ auto tOrP = make_tensor(rP.data(), convert_layout_acc_Aregs<typename T::Mma>(rP.layout()));
+ gemm_rs(acc_o, tOrP, rV, tOsV, mma, copyV, threadV);
+ };
+ step(std::true_type{}, block);
+ for (--block; block >= 0; --block)
+ step(std::false_type{}, block);
+ softmax.normalize(acc_o);
+ auto rO = to_half(acc_o);
+ auto sO = make_tensor(sQ.data(), typename T::QLayout{});
+ auto copyO = make_tiled_copy_C(Copy_Atom<AutoVectorizingCopyWithAssumedAlignment<128>, Element>{}, mma);
+ auto threadO = copyO.get_thread_slice(tid);
+ auto tOrO = threadO.retile_S(rO);
+ auto tOsO = threadO.partition_D(sO);
+ cute::copy(copyO, tOrO, tOsO);
+ 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<BM>, Int<D>>{}, 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<BM>, Int<D>>{}));
+ __syncthreads();
+ auto output_regs = make_tensor<Element>(shape(to_gO));
+ cute::copy(output_copy, from_sO, output_regs);
+ copy_rows<false, false, FullDim>(output_copy, output_regs, to_gO, output_id, p.nq - m_block * BM, p.dim);
+ }
+}
+
+static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+ a = index == 3 ? -a : a;
+ b = index == 2 ? -b : b;
+ c = index == 1 ? -c : c;
+ d = index == 0 ? -d : d;
+ return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+}
+
+static __global__ void quantize_attention(const half* src, int8_t* dst, float* scales, int64_t k, int64_t rows, float input_scale, float input_bias) {
+ const int64_t row = blockIdx.x;
+ const int tid = threadIdx.x;
+ if (row >= rows)
+ return;
+ extern __shared__ float values[];
+ __shared__ float maxima[32];
+ const half* src_row = src + row * k;
+ int8_t* dst_row = dst + row * k;
+ const int lane_id = tid % 32;
+ {
+ const int lane = lane_id;
+ constexpr int group_size = 256;
+ for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+ float v[8];
+#pragma unroll
+ for (int j = 0; j < 8; ++j) {
+ const float value = __fmaf_rn(__half2float(src_row[group * group_size + j * 32 + lane]), input_scale, input_bias);
+ v[j] = value * (1.0f / 16.0f);
+ }
+ // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+ for (int stride = 1; stride <= 4; stride *= 4) {
+ const int base = lane & ~(3 * stride);
+#pragma unroll
+ for (int j = 0; j < 8; ++j) {
+ const float a = __shfl_sync(0xffffffff, v[j], base, 32);
+ const float b = __shfl_sync(0xffffffff, v[j], base + stride, 32);
+ const float c = __shfl_sync(0xffffffff, v[j], base + 2 * stride, 32);
+ const float d = __shfl_sync(0xffffffff, v[j], base + 3 * stride, 32);
+ v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+ }
+ }
+#pragma unroll
+ for (int j = 0; j < 8; j += 2) {
+ const float a = __shfl_sync(0xffffffff, v[j], lane % 16, 32);
+ const float b = __shfl_sync(0xffffffff, v[j], lane % 16 + 16, 32);
+ const float c = __shfl_sync(0xffffffff, v[j + 1], lane % 16, 32);
+ const float d = __shfl_sync(0xffffffff, v[j + 1], lane % 16 + 16, 32);
+ v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
+ v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+ }
+#pragma unroll
+ for (int j = 0; j < 2; ++j) {
+ const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+ for (int t = 0; t < 4; ++t) {
+ values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+ }
+ }
+ }
+ __syncthreads();
+ }
+ float local_max = 0.0f;
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ local_max = fmaxf(local_max, fabsf(values[i]));
+ }
+ local_max = sd_warp_max(local_max);
+ const int lane = tid & 31;
+ const int warp = tid >> 5;
+ if (lane == 0) {
+ maxima[warp] = local_max;
+ }
+ __syncthreads();
+
+ if (warp == 0) {
+ local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+ local_max = sd_warp_max(local_max);
+ if (lane == 0) {
+ maxima[0] = local_max / 127.0f;
+ if (scales != nullptr) {
+ scales[row] = maxima[0];
+ } else {
+ *((float*)(dst_row + k)) = maxima[0];
+ }
+ }
+ }
+ __syncthreads();
+
+ const float row_scale = maxima[0];
+ const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+ value = max(-127, min(127, value));
+ dst_row[i] = (int8_t)value;
+ }
+}
+
+template <typename T>
+static __global__ void copy_to_half(const char* src, half* dst, int64_t count, int64_t dim, int64_t rows, int64_t heads, size_t nb1, size_t nb2, size_t nb3, float scale) {
+ const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (i < count) {
+ const int64_t row = i / dim;
+ const auto* p = reinterpret_cast<const T*>(src + (row % rows) * nb1 +
+ ((row / rows) % heads) * nb2 + (row / (rows * heads)) * nb3);
+ dst[i] = __float2half_rn(float(p[i % dim]) * scale);
+ }
+}
+
+static __global__ void copy_from_half(const half* src, float* dst, int64_t count, float scale, float bias, bool scaled) {
+ const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (i < count)
+ dst[i] = scaled ? __fmaf_rn(__half2float(src[i]), scale, bias) : __half2float(src[i]);
+}
+
+static float param_f32(const ggml_tensor* node, int index) {
+ float value;
+ std::memcpy(&value, reinterpret_cast<const char*>(node->op_params) + index * sizeof(value), sizeof(value));
+ return value;
+}
+
+static int padded_dim(int64_t dim) {
+ return dim == 80 || dim == 96 || dim == 112 ? int(dim) : int((dim + 63) / 64 * 64);
+}
+
+struct CUDADevice {
+ int ordinal;
+ cudaDeviceProp properties;
+};
+
+static const CUDADevice* cuda_device(ggml_backend_dev_t device) {
+ static std::mutex mutex;
+ static std::unordered_map<ggml_backend_dev_t, CUDADevice> devices;
+ std::lock_guard<std::mutex> lock(mutex);
+ auto found = devices.find(device);
+ if (found != devices.end())
+ return &found->second;
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+ if (!query)
+ return nullptr;
+ CUDADevice info = {};
+ info.ordinal = query(device);
+ if (info.ordinal < 0 || cudaGetDeviceProperties(&info.properties, info.ordinal) != cudaSuccess)
+ return nullptr;
+ return &devices.emplace(device, info).first->second;
+}
+
+static bool attention_supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+ static const bool disabled = std::getenv("GGML_CUDA_DISABLE_CUTE_FATTN") || std::getenv("GGML_CUDA_DISABLE_CUTLASS_FATTN");
+ const auto* info = cuda_device(device);
+ const auto* q = node->src[0];
+ const auto* k = node->src[1];
+ const auto* v = node->src[2];
+ if (disabled || !info || info->properties.major < 8 || !q || !k || !v ||
+ q->type != GGML_TYPE_F32 || node->type != GGML_TYPE_F32 ||
+ q->ne[0] < 32 || q->ne[0] > 256 || q->ne[0] % 8 ||
+ size_t(192 * padded_dim(q->ne[0]) * sizeof(half)) > info->properties.sharedMemPerBlockOptin ||
+ q->ne[1] < 1 || k->ne[1] < 1 || k->ne[2] < 1 ||
+ node->src[3] || node->src[4] || param_f32(node, 1) != 0.0f || param_f32(node, 2) != 0.0f ||
+ q->ne[2] % k->ne[2] || q->ne[3] != k->ne[3] || !ggml_are_same_shape(k, v) ||
+ q->nb[0] != sizeof(float) || !ggml_is_contiguous(node) ||
+ (k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_F32) ||
+ (v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_F32) ||
+ k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type))
+ return false;
+ return k->ne[0] == q->ne[0] && q->ne[1] <= INT_MAX / q->ne[0] && k->ne[1] <= INT_MAX / q->ne[0] &&
+ q->ne[2] <= 65535 && q->ne[3] <= 65535 && ggml_nelements(q) <= INT_MAX && ggml_nelements(k) <= INT_MAX;
+}
+
+static bool attention_preferred(ggml_backend_dev_t, const ggml_tensor* node, void*) {
+ const auto* q = node->src[0];
+ const auto* k = node->src[1];
+ if (q->ne[0] == 40 || q->ne[0] == 72)
+ return q->ne[1] >= 512 && k->ne[1] >= 512;
+ if (q->ne[0] == 80 || q->ne[0] == 96 || q->ne[0] == 112)
+ return false;
+ const int minimum = q->ne[0] == 128 ? 2048 : 4096;
+ return q->ne[1] >= minimum && k->ne[1] >= minimum;
+}
+
+static bool attention_selected(ggml_backend_dev_t device, const ggml_tensor* node) {
+ return node->op == GGML_OP_FLASH_ATTN_EXT && attention_supports(device, node, nullptr) &&
+ (attention_preferred(device, node, nullptr) || !ggml_ext_dev_supports_native_op(device, node));
+}
+
+static bool fusion_enabled() {
+ static const bool disabled = [] {
+ const char* value = std::getenv("GGML_CUDA_DISABLE_FUSION");
+ return value && std::atoi(value) != 0;
+ }();
+ return !disabled;
+}
+
+static const ggml_tensor* graph_node(const ggml_cgraph* graph, int index) {
+ return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
+}
+
+static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, int first) {
+ const auto* attn = graph_node(graph, first);
+ const auto* previous = attn;
+ const int total = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
+ bool scaled = false;
+ for (int count = 1; count < 8 && first + count < total; ++count) {
+ const auto* node = graph_node(graph, first + count);
+ int convrot_group;
+ if (node->src[0] != previous)
+ break;
+ if (node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE) {
+ if (node->view_offs || !ggml_is_contiguous(node) || ggml_nelements(node) != ggml_nelements(attn))
+ break;
+ } else if (node->op == GGML_OP_SCALE && !scaled) {
+ if (!ggml_is_contiguous(node) || !ggml_are_same_shape(previous, node))
+ break;
+ scaled = true;
+ } else if (sd_ggml_convrot_describe(node, &convrot_group)) {
+ if (convrot_group == 256 && previous->ne[0] % 256 == 0 &&
+ size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
+ return count + 1;
+ break;
+ } else
+ break;
+ previous = node;
+ }
+ if (first + 1 < total) {
+ const auto* next = graph_node(graph, first + 1);
+ if (next->op == GGML_OP_SCALE && next->src[0] == attn && ggml_is_contiguous(next) && ggml_are_same_shape(attn, next))
+ return 2;
+ }
+ return 1;
+}
+
+static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+ if (!fusion_enabled())
+ return 0;
+ int offset = 0;
+ if (ggml_ext_op_get(graph_node(graph, first)) == sd_ggml_i8_op()) {
+ if (first + 5 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
+ !sd_cuda_i8_rope_supports(device, graph_node(graph, first), graph_node(graph, first + 1),
+ graph_node(graph, first + 2), graph_node(graph, first + 3), graph_node(graph, first + 4)) ||
+ graph_node(graph, first + 5)->src[0] != graph_node(graph, first + 4))
+ return 0;
+ offset = 5;
+ } else if (graph_node(graph, first)->op == GGML_OP_RMS_NORM) {
+ if (first + 3 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) ||
+ !sd_cuda::rope_fusion_supports(graph_node(graph, first), graph_node(graph, first + 1), graph_node(graph, first + 2)) ||
+ graph_node(graph, first + 3)->src[0] != graph_node(graph, first + 2))
+ return 0;
+ offset = 3;
+ }
+ if (!attention_selected(device, graph_node(graph, first + offset)))
+ return 0;
+ const int count = offset + output_fusion(device, graph, first + offset);
+ return count > 1 ? count : 0;
+}
+
+static void attention_optimize(ggml_backend_dev_t device, ggml_cgraph* graph, void*) {
+ if (!fusion_enabled())
+ return;
+ for (int i = 3; i < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)); ++i) {
+ const auto* attn = graph_node(graph, i);
+ if (!attention_selected(device, attn))
+ continue;
+ for (int j = i - 3; j >= 0 && j >= i - 32; --j) {
+ if (attn->src[0] == graph_node(graph, j + 2) &&
+ sd_cuda::rope_fusion_supports(graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2))) {
+ if (j >= 2 && sd_cuda_i8_rope_supports(device, graph_node(graph, j - 2), graph_node(graph, j - 1), graph_node(graph, j), graph_node(graph, j + 1), graph_node(graph, j + 2)) &&
+ ggml_ext_graph_delay(graph, j - 2, 5, i))
+ break;
+ ggml_ext_graph_delay(graph, j, 3, i);
+ break;
+ }
+ }
+ }
+}
+
+struct AttentionPlan {
+ int index;
+ size_t offsets[4];
+};
+
+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 <int D, bool FullDim>
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+ auto kernel = sd_fattn_cute::kernel<64, 64, D, FullDim>;
+ constexpr int shared_bytes = sd_fattn_cute::Traits<64, 64, D>::SharedBytes;
+ static std::mutex mutex;
+ static std::unordered_map<int, bool> initialized;
+ {
+ std::lock_guard<std::mutex> lock(mutex);
+ if (!initialized[context.device]) {
+ auto error = cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shared_bytes);
+ if (error != cudaSuccess)
+ return error;
+ initialized[context.device] = true;
+ }
+ }
+ kernel<<<dim3((p.nq + 63) / 64, batches, p.heads), 128, shared_bytes, static_cast<cudaStream_t>(context.stream)>>>(p);
+ return cudaGetLastError();
+}
+
+template <int D>
+static cudaError_t attention_launch(const ggml_ext_cuda_context& context, const sd_fattn_cute::Params& p, int64_t batches) {
+ return p.dim == D ? attention_launch<D, true>(context, p, batches) : attention_launch<D, false>(context, p, batches);
+}
+
+static ggml_status attention_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+ const auto& plan = *static_cast<AttentionPlan*>(state);
+ const auto* attn = nodes[plan.index];
+ auto device = ggml_backend_get_device(execution->backend);
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+ ggml_ext_cuda_context context;
+ if (!query || !query(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<cudaStream_t>(context.stream);
+ auto* base = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ half* work[4];
+ for (int i = 0; i < 4; ++i)
+ work[i] = reinterpret_cast<half*>(base + plan.offsets[i]);
+ const half* inputs[3] = {work[0], work[1], work[2]};
+ for (int i = 0; i < 3; ++i) {
+ const auto* src = attn->src[i];
+ const float scale = i == 0 ? param_f32(attn, 0) : 1.0f;
+ if (i == 0 && plan.index == 5) {
+ if (sd_cuda_i8_rope_launch(stream, nodes, work[0], scale) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ } else if (i == 0 && plan.index == 3) {
+ if (sd_cuda::rms_rope_launch(stream, nodes, work[0], true, scale) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ } else if (i != 0 && src->type == GGML_TYPE_F16 && ggml_is_contiguous(src) && uintptr_t(src->data) % 16 == 0) {
+ inputs[i] = static_cast<const half*>(src->data);
+ } else {
+ auto kernel = src->type == GGML_TYPE_F16 ? copy_to_half<half> : copy_to_half<float>;
+ kernel<<<(ggml_nelements(src) + 255) / 256, 256, 0, stream>>>(static_cast<const char*>(src->data), work[i],
+ ggml_nelements(src), src->ne[0], src->ne[1], src->ne[2], src->nb[1], src->nb[2], src->nb[3], scale);
+ }
+ if (cudaGetLastError() != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ }
+ const auto* q = attn->src[0];
+ const auto* k = attn->src[1];
+ sd_fattn_cute::Params p = {reinterpret_cast<const cutlass::half_t*>(inputs[0]),
+ reinterpret_cast<const cutlass::half_t*>(inputs[1]), reinterpret_cast<const cutlass::half_t*>(inputs[2]),
+ reinterpret_cast<cutlass::half_t*>(work[3]), int(q->ne[1]), int(k->ne[1]), int(q->ne[2]), int(k->ne[2]), int(q->ne[0])};
+ 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<int8_t*>(dst->data);
+ if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ if (cudaFuncSetAttribute(quantize_attention, cudaFuncAttributeMaxDynamicSharedMemorySize,
+ int(cuda_device(device)->properties.sharedMemPerBlockOptin - 128)) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ quantize_attention<<<rows, 256, width * sizeof(float), stream>>>(work[3], packed,
+ reinterpret_cast<float*>(packed + width * padded_rows), width, rows, scale, bias);
+ } else {
+ copy_from_half<<<(ggml_nelements(attn) + 255) / 256, 256, 0, stream>>>(work[3], static_cast<float*>(dst->data), ggml_nelements(attn), scale, bias, dst->op == GGML_OP_SCALE);
+ }
+ return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+static void attention_destroy(void* state, void*) {
+ delete static_cast<AttentionPlan*>(state);
+}
+
+void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
+ if (!cuda_device(device))
+ return;
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.cute_attention.cuda.v1";
+ desc.op = GGML_OP_FLASH_ATTN_EXT;
+ desc.priority = 200;
+ desc.supports = attention_supports;
+ desc.preferred = attention_preferred;
+ desc.match = attention_match;
+ desc.optimize = attention_optimize;
+ desc.prepare = attention_prepare;
+ desc.execute = attention_execute;
+ desc.destroy = attention_destroy;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.rms_rope_cute_attention.cuda.v1";
+ desc.op = GGML_OP_RMS_NORM;
+ desc.priority = 300;
+ desc.supports = nullptr;
+ desc.preferred = nullptr;
+ desc.optimize = nullptr;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.i8_rope_cute_attention.cuda.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_i8_op();
+ desc.priority = 400;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}
diff --git a/src/core/cuda/i8_gated_residual.cu b/src/core/cuda/i8_gated_residual.cu
new file mode 100644
index 0000000..21b6134
--- /dev/null
+++ b/src/core/cuda/i8_gated_residual.cu
@@ -0,0 +1,461 @@
+#include <cuda_runtime.h>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <unordered_map>
+#include "core/ggml_extensions.h"
+
+namespace sd_i8_gated {
+ struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+ using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+ using Base::Base;
+
+ CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+ const auto shape = tiled_shape();
+ if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+ // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+ constexpr int band_width = 16;
+ if (index >= shape.m() * shape.n()) {
+ return {shape.m(), shape.n(), get_batch_idx()};
+ }
+ const int band = index / (shape.m() * band_width);
+ const int width = min(band_width, shape.n() - band * band_width);
+ const int within = index - band * shape.m() * band_width;
+ return {within / width, band * band_width + within % width, get_batch_idx()};
+ }
+#endif
+ return Base::get_tile_offset(index);
+ }
+ };
+
+ template <int M, int N, int WM>
+ struct TileConfig {
+ using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+ using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+ };
+
+ enum Tile {
+ TILE_128_256,
+ TILE_128_128,
+ TILE_64_128,
+ TILE_16_128,
+ };
+
+ static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+ static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+ return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+ }
+
+ static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+ if (rows < 32) {
+ return TILE_16_128;
+ }
+ const int64_t columns = (n + 127) / 128;
+ if (((rows + 15) / 16) * columns <= nsm) {
+ return TILE_16_128;
+ }
+ if (((rows + 63) / 64) * columns <= nsm) {
+ return TILE_64_128;
+ }
+ const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
+ const int64_t square_tiles = ((rows + 127) / 128) * columns;
+ // A narrower tile costs more per output element, but can waste less of the last wave.
+ const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+ const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+ return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+ }
+
+ template <class T>
+ struct MulRound {
+ CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
+ T result;
+ CUTLASS_PRAGMA_UNROLL
+ for (int i = 0; i < T::kElements; ++i) {
+ result[i] = __fmul_rn(a[i], b[i]);
+ }
+ return result;
+ }
+ };
+
+ template <class T>
+ struct AddRound {
+ CUTLASS_DEVICE T operator()(const T& a, const T& b) const {
+ T result;
+ CUTLASS_PRAGMA_UNROLL
+ for (int i = 0; i < T::kElements; ++i) {
+ result[i] = __fadd_rn(a[i], b[i]);
+ }
+ return result;
+ }
+ };
+
+ template <class Config, bool scaled, bool stream_k>
+ struct KernelTypes {
+ using Tile = typename Config::Tile;
+ using Warp = typename Config::Warp;
+ using Swizzle = std::conditional_t<stream_k, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+ using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+ using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+ using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+ using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+ using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+ using Residual = cutlass::epilogue::threadblock::VisitorAuxLoad<Map, float, cute::Stride<int64_t, cute::_1, int64_t>>;
+ using GateMul = cutlass::epilogue::threadblock::VisitorCompute<MulRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+ using Add = cutlass::epilogue::threadblock::VisitorCompute<AddRound, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+ using Scalar = cutlass::epilogue::threadblock::VisitorScalarBroadcast<float>;
+ using ScaleOut = cutlass::epilogue::threadblock::Sm80EVT<GateMul, ScaleW, Scalar>;
+ using Scaled = std::conditional_t<scaled, ScaleOut, ScaleW>;
+ using Gated = cutlass::epilogue::threadblock::Sm80EVT<GateMul, Scaled, Row>;
+ using WithResidual = cutlass::epilogue::threadblock::Sm80EVT<Add, Gated, Residual>;
+ using Output = WithResidual;
+ using Store = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+ using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, Output>;
+ using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+ int8_t,
+ cutlass::layout::RowMajor,
+ cutlass::ComplexTransform::kNone,
+ 16,
+ int8_t,
+ cutlass::layout::ColumnMajor,
+ cutlass::ComplexTransform::kNone,
+ 16,
+ float,
+ cutlass::layout::RowMajor,
+ 4,
+ int,
+ float,
+ cutlass::arch::OpClassTensorOp,
+ cutlass::arch::Sm80,
+ Tile,
+ Warp,
+ cutlass::gemm::GemmShape<16, 8, 32>,
+ Epilogue,
+ Swizzle,
+ 4,
+ cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+ };
+
+ struct DeviceInfo {
+ int ordinal = -1;
+ int major = 0;
+ int sms = 0;
+ int threads_per_sm = 0;
+ size_t shared = 0;
+ };
+ static DeviceInfo device_info(ggml_backend_dev_t device) {
+ static std::mutex mutex;
+ static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
+ std::lock_guard<std::mutex> lock(mutex);
+ auto found = cache.find(device);
+ if (found != cache.end())
+ return found->second;
+ DeviceInfo info;
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+ int ordinal = query ? query(device) : -1;
+ cudaDeviceProp props = {};
+ if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
+ info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+ cache.emplace(device, info);
+ return info;
+ }
+ static float param_f32(const ggml_tensor* tensor, int index) {
+ float value;
+ std::memcpy(&value, tensor->op_params + index, sizeof(value));
+ return value;
+ }
+ static bool aligned(const ggml_tensor* tensor) {
+ if (!tensor)
+ return false;
+ size_t offset = 0;
+ while (tensor->view_src) {
+ offset += tensor->view_offs;
+ tensor = tensor->view_src;
+ }
+ return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
+ }
+ struct Operation {
+ SDI8Operation mat;
+ const ggml_tensor* gate;
+ const ggml_tensor* residual;
+ const ggml_tensor* prefix;
+ float scale;
+ int64_t rows, n, k;
+ };
+ static Operation describe(ggml_tensor* const* nodes, int count) {
+ Operation op = {};
+ sd_ggml_i8_describe(nodes[0], &op.mat);
+ const bool concat = nodes[count - 1]->op == GGML_OP_CONCAT;
+ const auto* add = nodes[count - (concat ? 2 : 1)];
+ const auto* mul = nodes[count - (concat ? 3 : 2)];
+ op.gate = mul->src[1];
+ op.residual = add->src[0] == mul ? add->src[1] : add->src[0];
+ op.prefix = concat ? nodes[count - 1]->src[0] : nullptr;
+ op.scale = nodes[1]->op == GGML_OP_SCALE ? param_f32(nodes[1], 0) : 1.f;
+ op.rows = ggml_nrows(nodes[0]);
+ op.n = nodes[0]->ne[0];
+ op.k = op.mat.weight->ne[0];
+ return op;
+ }
+ template <class Types>
+ static typename Types::Gemm::Arguments arguments(const Operation& op, float* output, bool metadata) {
+ const auto* x = metadata ? nullptr : static_cast<const int8_t*>(op.mat.input->data);
+ const auto* w = metadata ? nullptr : static_cast<const int8_t*>(op.mat.weight->data);
+ const auto* xs = metadata ? nullptr : op.mat.packed ? reinterpret_cast<const float*>(x + op.k * ((op.rows + 3) / 4 * 4))
+ : static_cast<const float*>(op.mat.input_scale->data);
+ const auto* ws = metadata ? nullptr : static_cast<const float*>(op.mat.weight_scale->data);
+ const auto* gate = metadata ? nullptr : static_cast<const float*>(op.gate->data);
+ auto* residual = metadata ? nullptr : static_cast<float*>(op.residual->data);
+ typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+ typename Types::Scaled::Arguments scaled;
+ if constexpr (std::is_same<typename Types::Scaled, typename Types::ScaleW>::value)
+ scaled = scales;
+ else
+ scaled = {scales, {{op.scale}}, {}};
+ typename Types::Epilogue::Arguments epilogue = {{{scaled, {gate, 0.f, {}}, {}}, {residual, 0.f, {op.n, cute::_1{}, op.rows * op.n}}, {}}, {output, {op.n, cute::_1{}, op.rows * op.n}}};
+ return {cutlass::gemm::GemmUniversalMode::kGemm, {int(op.rows), int(op.n), int(op.k)}, 1, epilogue, x, w, nullptr, nullptr, op.rows * op.k, op.n * op.k, 0, 0, op.k, op.k, 0, 0};
+ }
+ struct Plan {
+ DeviceInfo device;
+ size_t kernel_workspace = 0;
+ size_t output_offset = 0;
+ int occupancy = 0;
+ ggml_status (*launch)(Plan&, const Operation&, float*, void*, cudaStream_t) = nullptr;
+ };
+ template <class Types>
+ static ggml_status launch(Plan& plan, const Operation& op, float* output, void* scratch, cudaStream_t stream) {
+ using Kernel = typename Types::Kernel;
+ auto args = arguments<Types>(op, output, false);
+ if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+ return GGML_STATUS_FAILED;
+ if (plan.occupancy == 0) {
+ if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
+ cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ int occupancy = 0;
+ if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&occupancy, cutlass::Kernel2<Kernel>,
+ Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
+ occupancy <= 0)
+ return GGML_STATUS_FAILED;
+ plan.occupancy = occupancy;
+ }
+ typename Kernel::Params params(args, plan.device.sms, plan.occupancy);
+ if (params.get_workspace_size() > plan.kernel_workspace || params.init_workspace(scratch, stream) != cutlass::Status::kSuccess)
+ return GGML_STATUS_FAILED;
+ cutlass::arch::synclog_setup();
+ cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
+ return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+ template <class Config, bool Scaled, bool StreamK>
+ static bool configure(const Operation& op, Plan& plan) {
+ using Types = KernelTypes<Config, Scaled, StreamK>;
+ using Kernel = typename Types::Kernel;
+ if (sizeof(typename Kernel::SharedStorage) > plan.device.shared)
+ return false;
+ auto args = arguments<Types>(op, nullptr, true);
+ if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+ return false;
+ size_t workspace = 0;
+ // Bound every possible occupancy without setting a CUDA device or initializing a kernel during graph measurement.
+ for (int occupancy = 1; occupancy <= plan.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+ typename Kernel::Params params(args, plan.device.sms, occupancy);
+ const auto grid = params.get_grid_dims();
+ if (grid.y > 65535 || grid.z > 65535)
+ return false;
+ workspace = std::max(workspace, params.get_workspace_size());
+ }
+ plan.kernel_workspace = workspace;
+ plan.output_offset = (workspace + 255) / 256 * 256;
+ plan.launch = launch<Types>;
+ return true;
+ }
+ template <class Config>
+ static bool configure(const Operation& op, Plan& plan) {
+ if (op.scale != 1.f)
+ return op.rows >= 8192 ? configure<Config, true, true>(op, plan) : configure<Config, true, false>(op, plan);
+ return op.rows >= 8192 ? configure<Config, false, true>(op, plan) : configure<Config, false, false>(op, plan);
+ }
+ static bool configure(const Operation& op, Plan& plan) {
+ auto try_tile = [&](Tile tile) {
+ switch (tile) {
+ case TILE_128_256:
+ return configure<TileConfig<128, 256, 64>>(op, plan);
+ case TILE_128_128:
+ return configure<TileConfig<128, 128, 64>>(op, plan);
+ case TILE_64_128:
+ return configure<TileConfig<64, 128, 32>>(op, plan);
+ case TILE_16_128:
+ return configure<TileConfig<16, 128, 16>>(op, plan);
+ }
+ return false;
+ };
+ const auto tile = select_tile(op.rows, op.n, plan.device.sms);
+ if (try_tile(tile))
+ return true;
+ for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
+ if (fallback != tile && try_tile(fallback))
+ return true;
+ return false;
+ }
+ static bool fusion_disabled() {
+ static const bool disabled = [] { const char* v = std::getenv("GGML_CUDA_DISABLE_FUSION"); return v && std::atoi(v) != 0; }();
+ return disabled;
+ }
+ static void optimize(ggml_backend_dev_t, ggml_cgraph* graph, void*) {
+#ifndef SD_CUDA_FORCE_CUBLAS
+ if (fusion_disabled())
+ return;
+ const int total = ggml_graph_n_nodes(graph);
+ for (int i = 2; i + 1 < total; ++i) {
+ const auto* mul = ggml_graph_node(graph, i);
+ const auto* add = ggml_graph_node(graph, i + 1);
+ if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || (add->src[0] != mul && add->src[1] != mul))
+ continue;
+ const auto* producer = mul->src[0];
+ const bool scaled = producer->op == GGML_OP_SCALE;
+ const auto* mat = scaled ? producer->src[0] : producer;
+ SDI8Operation op;
+ if (!sd_ggml_i8_describe(mat, &op) || op.input->type != GGML_TYPE_I8 || !op.weight_scale || op.bias)
+ continue;
+ int scale_index = -1;
+ for (int j = i - 1; j >= 0 && j >= i - 64; --j) {
+ if (scaled && ggml_graph_node(graph, j) == producer)
+ scale_index = j;
+ if (ggml_graph_node(graph, j) != mat)
+ continue;
+ if (scaled) {
+ if (scale_index < 0 || (j + 1 < scale_index && !ggml_ext_graph_delay(graph, j, 1, scale_index)))
+ break;
+ j = scale_index - 1;
+ }
+ ggml_ext_graph_delay(graph, j, scaled ? 2 : 1, i);
+ break;
+ }
+ }
+#endif
+ }
+ static int match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+ return 0;
+#else
+ const int available = ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)) - first;
+ if (fusion_disabled() || available < 3)
+ return 0;
+ ggml_tensor* nodes[5] = {};
+ for (int i = 0; i < std::min(available, 5); ++i)
+ nodes[i] = ggml_graph_node(const_cast<ggml_cgraph*>(graph), first + i);
+ SDI8Operation mat;
+ if (!sd_ggml_i8_describe(nodes[0], &mat) || mat.input->type != GGML_TYPE_I8 || !mat.weight_scale || mat.bias)
+ return 0;
+ const int64_t rows = ggml_nrows(nodes[0]), n = nodes[0]->ne[0], k = mat.weight->ne[0];
+ if (rows < (small_rows_enabled(rows, n, k) ? 1 : 128) || rows > INT_MAX || n > INT_MAX || k > INT_MAX || n % 128 || k % 64 ||
+ ((rows + 15) / 16) > INT_MAX / ((n + 127) / 128))
+ return 0;
+ for (const auto* t : {mat.weight, mat.input, mat.weight_scale, mat.packed ? mat.weight_scale : mat.input_scale})
+ if (!aligned(t))
+ return 0;
+ const bool scaled = nodes[1]->op == GGML_OP_SCALE;
+ int count = scaled ? 4 : 3;
+ if (available < count || (scaled && (nodes[1]->src[0] != nodes[0] || param_f32(nodes[1], 1) != 0.f)))
+ return 0;
+ const auto* mul = nodes[count - 2];
+ const auto* add = nodes[count - 1];
+ if (mul->op != GGML_OP_MUL || add->op != GGML_OP_ADD || mul->src[0] != nodes[count - 3] ||
+ (add->src[0] != mul && add->src[1] != mul))
+ return 0;
+ const auto* gate = mul->src[1];
+ const auto* residual = add->src[0] == mul ? add->src[1] : add->src[0];
+ if (!gate || !residual || gate->ne[0] != n || ggml_nelements(gate) != n || !aligned(gate) || !aligned(residual))
+ return 0;
+ for (const auto* t : {gate, residual, mul, add})
+ if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || (t != gate && !ggml_are_same_shape(nodes[0], t)))
+ return 0;
+ if (available > count) {
+ const auto* concat = nodes[count];
+ if (concat->op == GGML_OP_CONCAT && concat->src[1] == add && concat->op_params[0] == 1 &&
+ concat->type == GGML_TYPE_F32 && concat->ne[2] == 1 && concat->ne[3] == 1 && ggml_is_contiguous(concat) &&
+ concat->src[0]->type == GGML_TYPE_F32 && ggml_is_contiguous(concat->src[0]))
+ ++count;
+ }
+ const auto* prefix = nodes[count - 1]->op == GGML_OP_CONCAT ? nodes[count - 1]->src[0] : nullptr;
+ for (int i = 0; i < count; ++i)
+ if (nodes[i] == gate || nodes[i] == residual || nodes[i] == prefix)
+ return 0;
+ Plan plan;
+ plan.device = device_info(device);
+ if (plan.device.major < 8 || plan.device.sms <= 0)
+ return 0;
+ return configure(describe(nodes, count), plan) ? count : 0;
+#endif
+ }
+ static 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<Plan*>(state);
+ }
+ static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ if (!b)
+ return false;
+ const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+ }
+ static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+ auto& plan = *static_cast<Plan*>(state);
+ const auto op = describe(nodes, count);
+ auto* dst = nodes[count - 1];
+ auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+ auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+ ggml_ext_cuda_context context;
+ if (!query || !query(execution->backend, &context) || context.device != plan.device.ordinal)
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<cudaStream_t>(context.stream);
+ auto* scratch = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ bool alias = reinterpret_cast<uintptr_t>(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<char*>(dst->data);
+ const size_t prefix_bytes = op.prefix ? ggml_nbytes(op.prefix) : 0;
+ auto status = plan.launch(plan, op, reinterpret_cast<float*>(output + prefix_bytes), scratch, stream);
+ if (status != GGML_STATUS_SUCCESS)
+ return status;
+ if (prefix_bytes && cudaMemcpyAsync(output, op.prefix->data, prefix_bytes, cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ if (alias && cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ return GGML_STATUS_SUCCESS;
+ }
+}
+
+void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.i8_gated_residual.cuda.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_i8_op();
+ desc.priority = 325;
+ desc.optimize = sd_i8_gated::optimize;
+ desc.match = sd_i8_gated::match;
+ desc.prepare = sd_i8_gated::prepare;
+ desc.execute = sd_i8_gated::execute;
+ desc.destroy = sd_i8_gated::destroy;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}
diff --git a/src/core/cuda/i8_gemm.cu b/src/core/cuda/i8_gemm.cu
new file mode 100644
index 0000000..69f91d2
--- /dev/null
+++ b/src/core/cuda/i8_gemm.cu
@@ -0,0 +1,330 @@
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <cutlass/gemm/threadblock/threadblock_swizzle_streamk.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <unordered_map>
+#include "i8_gemm.h"
+
+namespace sd_i8_cutlass {
+ struct BandedStreamK : cutlass::gemm::threadblock::ThreadblockSwizzleStreamK {
+ using Base = cutlass::gemm::threadblock::ThreadblockSwizzleStreamK;
+ using Base::Base;
+
+ CUTLASS_DEVICE cutlass::gemm::GemmCoord get_tile_offset(int index) const {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 890
+ const auto shape = tiled_shape();
+ if (problem_size.m() >= 4096 && problem_size.n() >= 16384 && shape.m() >= shape.n()) {
+ // Column bands retain weight tiles in Ada's L2 across successive row tiles.
+ constexpr int band_width = 16;
+ if (index >= shape.m() * shape.n()) {
+ return {shape.m(), shape.n(), get_batch_idx()};
+ }
+ const int band = index / (shape.m() * band_width);
+ const int width = min(band_width, shape.n() - band * band_width);
+ const int within = index - band * shape.m() * band_width;
+ return {within / width, band * band_width + within % width, get_batch_idx()};
+ }
+#endif
+ return Base::get_tile_offset(index);
+ }
+ };
+
+ template <int M, int N, int WM>
+ struct TileConfig {
+ using Tile = cutlass::gemm::GemmShape<M, N, 64>;
+ using Warp = cutlass::gemm::GemmShape<WM, 64, 64>;
+ };
+
+ enum Tile {
+ TILE_128_256,
+ TILE_128_128,
+ TILE_64_128,
+ TILE_16_128,
+ };
+
+ static bool small_rows_enabled(int64_t rows, int64_t n, int64_t k) {
+ static const char* setting = std::getenv("GGML_CUDA_I8_SMALL_ROWS");
+ return setting ? std::atoi(setting) != 0 : rows < 32 && n >= 4096 && k >= 4096;
+ }
+
+ static Tile select_tile(int64_t rows, int64_t n, int nsm) {
+ if (rows < 32) {
+ return TILE_16_128;
+ }
+ const int64_t columns = (n + 127) / 128;
+ if (((rows + 15) / 16) * columns <= nsm) {
+ return TILE_16_128;
+ }
+ if (((rows + 63) / 64) * columns <= nsm) {
+ return TILE_64_128;
+ }
+ const int64_t wide_tiles = ((rows + 127) / 128) * ((n + 255) / 256);
+ const int64_t square_tiles = ((rows + 127) / 128) * columns;
+ // A narrower tile costs more per output element, but can waste less of the last wave.
+ const double wide_cost = double((wide_tiles + nsm - 1) / nsm) / wide_tiles;
+ const double square_cost = 1.08 * double((square_tiles + nsm - 1) / nsm) / square_tiles;
+ return wide_cost <= square_cost ? TILE_128_256 : TILE_128_128;
+ }
+
+ struct DeviceInfo {
+ int ordinal = -1;
+ int major = 0;
+ int sms = 0;
+ int threads_per_sm = 0;
+ size_t shared = 0;
+ };
+ static DeviceInfo device_info(ggml_backend_dev_t device) {
+ static std::mutex mutex;
+ static std::unordered_map<ggml_backend_dev_t, DeviceInfo> cache;
+ std::lock_guard<std::mutex> lock(mutex);
+ auto found = cache.find(device);
+ if (found != cache.end())
+ return found->second;
+ DeviceInfo info;
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+ int ordinal = query ? query(device) : -1;
+ cudaDeviceProp props = {};
+ if (ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess)
+ info = {ordinal, props.major, props.multiProcessorCount, props.maxThreadsPerMultiProcessor, size_t(props.sharedMemPerBlockOptin)};
+ cache.emplace(device, info);
+ return info;
+ }
+ // The iterator pairs gate/up columns while reading the original checkpoint layout.
+ struct WeightLayout : cutlass::layout::PermuteBase {
+ int stride, half_n;
+ CUTLASS_HOST_DEVICE WeightLayout(cutlass::PitchLinearCoord extent, int64_t ld)
+ : stride(static_cast<int>(ld)), half_n(extent.strided() / 2) {}
+ CUTLASS_HOST_DEVICE WeightLayout(cutlass::MatrixCoord extent, int64_t ld)
+ : stride(static_cast<int>(ld)), half_n(extent.column() / 2) {}
+ CUTLASS_HOST_DEVICE int64_t operator()(cutlass::PitchLinearCoord c) const { return c.contiguous() + int64_t(c.strided() / 2 + (c.strided() % 2) * half_n) * stride; }
+ CUTLASS_HOST_DEVICE int64_t operator()(cutlass::MatrixCoord c) const { return c.row() + int64_t(c.column() / 2 + (c.column() % 2) * half_n) * stride; }
+ };
+
+ template <class Map>
+ struct SwigluStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+ using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+ using Base::Base;
+ template <class CB>
+ struct Callbacks : CB {
+ CUTLASS_DEVICE Callbacks(CB const& base)
+ : CB(base) {}
+ CUTLASS_DEVICE void end_step(int step) {
+ auto src = cute::filter(this->tC_rAux);
+ auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
+ CUTLASS_PRAGMA_UNROLL
+ for (int i = 0; i < cute::size(src); ++i) {
+ auto coord = coords(i);
+ if (cute::elem_less(coord, this->problem_shape)) {
+ // Each vector contains two gate/up pairs; only their SwiGLU results are stored.
+ const float4 a = *reinterpret_cast<const float4*>(&src(i));
+ float2 value = make_float2(__fmul_rn(a.x / (1.f + expf(-a.x)), a.y), __fmul_rn(a.z / (1.f + expf(-a.z)), a.w));
+ int64_t n = cute::get<1>(this->problem_shape), m = cute::get<0>(this->problem_shape);
+ int64_t offset = cute::get<0>(coord) * (n / 2) + cute::get<1>(coord) / 2 + cute::get<2>(coord) * m * (n / 2);
+ *reinterpret_cast<float2*>(this->params_ptr->ptr_aux + offset) = value;
+ }
+ }
+ }
+ };
+ template <class Shape>
+ CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
+ auto cb = Base::get_callbacks(tile, tid, shape);
+ return Callbacks<decltype(cb)>(cb);
+ }
+ };
+
+ static __global__ void pair_scales(const float* src, 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 <class Config, bool StreamK>
+ struct DenseTypes {
+ using Tile = typename Config::Tile;
+ using Warp = typename Config::Warp;
+ using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+ using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+ using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+ using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+ using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+ using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+ using Store = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+ using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+ using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+ int8_t,
+ cutlass::layout::RowMajor,
+ cutlass::ComplexTransform::kNone,
+ 16,
+ int8_t,
+ cutlass::layout::ColumnMajor,
+ cutlass::ComplexTransform::kNone,
+ 16,
+ float,
+ cutlass::layout::RowMajor,
+ 4,
+ int,
+ float,
+ cutlass::arch::OpClassTensorOp,
+ cutlass::arch::Sm80,
+ Tile,
+ Warp,
+ cutlass::gemm::GemmShape<16, 8, 32>,
+ Epilogue,
+ Swizzle,
+ 4,
+ cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+ };
+ template <class Config, bool StreamK>
+ struct SwigluTypes {
+ using Tile = typename Config::Tile;
+ using Warp = typename Config::Warp;
+ using Swizzle = std::conditional_t<StreamK, BandedStreamK, cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>>;
+ using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+ using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+ using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+ using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+ using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+ using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+ using Store = SwigluStore<Map>;
+ using Epilogue = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+ using GemmBase = typename cutlass::gemm::kernel::DefaultGemmUniversal<int8_t, cutlass::layout::RowMajor, cutlass::ComplexTransform::kNone, 16, int8_t, cutlass::layout::ColumnMajor, cutlass::ComplexTransform::kNone, 16, float, cutlass::layout::RowMajor, int, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80, Tile, Warp, cutlass::gemm::GemmShape<16, 8, 32>, cutlass::epilogue::thread::LinearCombination<float, 4, int, float>, Swizzle, 4, cutlass::arch::OpMultiplyAddSaturate, cutlass::gemm::SharedMemoryClearOption::kNone, false, false, false, cutlass::layout::NoPermute, cutlass::layout::NoPermute, WeightLayout>::GemmKernel;
+ using VisitorEpilogue = cutlass::epilogue::threadblock::EpilogueWithVisitorCallbacks<typename GemmBase::Epilogue, Epilogue, 1>;
+ using Kernel = std::conditional_t<StreamK,
+ cutlass::gemm::kernel::GemmWithEpilogueVisitorStreamk<typename GemmBase::Mma, VisitorEpilogue, Swizzle>,
+ cutlass::gemm::kernel::GemmWithEpilogueVisitor<typename GemmBase::Mma, VisitorEpilogue, Swizzle>>;
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+ };
+ 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 <class Types>
+ static typename Types::Gemm::Arguments arguments(const Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out) {
+ typename Types::ScaleW::Arguments scales = {{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}};
+ typename Types::Epilogue::Arguments epilogue = {scales, {out, {p.n, cute::_1{}, p.rows * p.n}}};
+ return {cutlass::gemm::GemmUniversalMode::kGemm, {int(p.rows), int(p.n), int(p.k)}, 1, epilogue, x, w, nullptr, nullptr, p.rows * p.k, p.n * p.k, 0, 0, p.k, p.k, 0, 0};
+ }
+ template <class Types>
+ static ggml_status launch(Plan& p, const int8_t* x, const int8_t* w, const float* xs, const float* ws, float* out, void* scratch, cudaStream_t stream) {
+ using Kernel = typename Types::Kernel;
+ if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
+ reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(out)) %
+ 16)
+ return GGML_STATUS_FAILED;
+ auto* paired = static_cast<float*>(scratch);
+ auto args = arguments<Types>(p, x, w, xs, p.swiglu ? paired : ws, out);
+ if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+ return GGML_STATUS_FAILED;
+ if (!p.occupancy) {
+ if (sizeof(typename Kernel::SharedStorage) >= (48 << 10) &&
+ cudaFuncSetAttribute(cutlass::Kernel2<Kernel>, cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(typename Kernel::SharedStorage)) != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ if (cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags(&p.occupancy, cutlass::Kernel2<Kernel>, Kernel::kThreadCount,
+ sizeof(typename Kernel::SharedStorage), cudaOccupancyDisableCachingOverride) != cudaSuccess ||
+ p.occupancy <= 0)
+ return GGML_STATUS_FAILED;
+ }
+ typename Kernel::Params params(args, p.device.sms, p.occupancy);
+ if (params.get_workspace_size() > p.kernel_bytes)
+ return GGML_STATUS_FAILED;
+ if (p.swiglu) {
+ pair_scales<<<unsigned((p.n + 255) / 256), 256, 0, stream>>>(ws, paired, int(p.n));
+ if (cudaGetLastError() != cudaSuccess)
+ return GGML_STATUS_FAILED;
+ }
+ if (params.init_workspace(static_cast<char*>(scratch) + p.scratch_offset, stream) != cutlass::Status::kSuccess)
+ return GGML_STATUS_FAILED;
+ cutlass::arch::synclog_setup();
+ cutlass::Kernel2<Kernel><<<params.get_grid_dims(), Kernel::kThreadCount, sizeof(typename Kernel::SharedStorage), stream>>>(params);
+ return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+ template <class Types>
+ static bool configure_kernel(Plan& p) {
+ using Kernel = typename Types::Kernel;
+ if (sizeof(typename Kernel::SharedStorage) > p.device.shared)
+ return false;
+ auto args = arguments<Types>(p, nullptr, nullptr, nullptr, nullptr, nullptr);
+ if (Kernel::can_implement(args) != cutlass::Status::kSuccess)
+ return false;
+ size_t bytes = 0;
+ for (int occupancy = 1; occupancy <= p.device.threads_per_sm / Kernel::kThreadCount; ++occupancy) {
+ typename Kernel::Params params(args, p.device.sms, occupancy);
+ const auto grid = params.get_grid_dims();
+ if (grid.y > 65535 || grid.z > 65535)
+ return false;
+ bytes = std::max(bytes, params.get_workspace_size());
+ }
+ p.kernel_bytes = bytes;
+ p.scratch_offset = p.swiglu ? (size_t(p.n) * sizeof(float) + 255) / 256 * 256 : 0;
+ p.launch = launch<Types>;
+ return true;
+ }
+ template <class Config>
+ static bool configure_tile(Plan& p) {
+ if (p.swiglu)
+ return p.rows >= 8192 ? configure_kernel<SwigluTypes<Config, true>>(p) : configure_kernel<SwigluTypes<Config, false>>(p);
+ return p.rows >= 8192 ? configure_kernel<DenseTypes<Config, true>>(p) : configure_kernel<DenseTypes<Config, false>>(p);
+ }
+ static bool configure(Plan& p) {
+ auto try_tile = [&](Tile tile) {
+ switch (tile) {
+ case TILE_128_256:
+ return configure_tile<TileConfig<128, 256, 64>>(p);
+ case TILE_128_128:
+ return configure_tile<TileConfig<128, 128, 64>>(p);
+ case TILE_64_128:
+ return configure_tile<TileConfig<64, 128, 32>>(p);
+ case TILE_16_128:
+ return configure_tile<TileConfig<16, 128, 16>>(p);
+ }
+ return false;
+ };
+ const auto tile = select_tile(p.rows, p.n, p.device.sms);
+ if (try_tile(tile))
+ return true;
+ for (auto fallback : {TILE_128_256, TILE_128_128, TILE_64_128, TILE_16_128})
+ if (fallback != tile && try_tile(fallback))
+ return true;
+ return false;
+ }
+}
+struct SDI8CudaGemmPlan : sd_i8_cutlass::Plan {};
+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 0000000..189df6e
--- /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 <cuda_runtime.h>
+#include "ggml-ext.h"
+struct SDI8CudaGemmPlan;
+SDI8CudaGemmPlan* sd_cuda_i8_gemm_prepare(ggml_backend_dev_t device, int64_t rows, int64_t n, int64_t k, bool swiglu, size_t* workspace);
+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 0000000..2f8ec3f
--- /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 <cuda_fp16.h>
+#include <cutlass/gemm/device/gemm_universal_adapter.h>
+#include <cutlass/gemm/kernel/default_gemm_universal_with_visitor.h>
+#include <algorithm>
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <cutlass/epilogue/threadblock/fusion/visitors.hpp>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+
+static __device__ __forceinline__ float warp_sum8(float value) {
+#pragma unroll
+ for (int offset = 4; offset > 0; offset >>= 1)
+ value = __fadd_rn(value, __shfl_xor_sync(0xffffffff, value, offset, 8));
+ return value;
+}
+
+template <class Map, class Out, bool rms = true>
+struct I8RopeStore : cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>> {
+ using Base = cutlass::epilogue::threadblock::VisitorAuxStore<Map, float, cutlass::FloatRoundStyle::round_to_nearest, cute::Stride<int64_t, cute::_1, int64_t>>;
+ struct Arguments : Base::Arguments {
+ const float* weight;
+ const float* pe;
+ float eps;
+ float scale;
+ int64_t tokens;
+ half* half_output;
+ int64_t prefix_tokens;
+ };
+ using Params = Arguments;
+ using SharedStorage = typename Base::SharedStorage;
+ const Params* extra;
+ template <class Shape>
+ static constexpr Params to_underlying_arguments(Shape const&, Arguments const& args, void*) { return args; }
+ CUTLASS_HOST_DEVICE I8RopeStore(Params const& p, SharedStorage const& s)
+ : Base(p, s), extra(&p) {}
+ template <class CB>
+ struct Callbacks : CB {
+ const Params* extra;
+ CUTLASS_DEVICE Callbacks(CB const& base, const Params* p)
+ : CB(base), extra(p) {}
+ CUTLASS_DEVICE float sum32(float4 v) {
+ float a = warp_sum8(__fmul_rn(v.x, v.x));
+ float b = warp_sum8(__fmul_rn(v.y, v.y));
+ float c = warp_sum8(__fmul_rn(v.z, v.z));
+ float d = warp_sum8(__fmul_rn(v.w, v.w));
+ return __fadd_rn(__fadd_rn(a, c), __fadd_rn(b, d));
+ }
+ CUTLASS_DEVICE void end_step(int step) {
+ static_assert(Map::Detail::kAccessWidth == 16);
+ auto src = cute::filter(this->tC_rAux);
+ auto coords = cute::filter(this->tC_cAux(cute::_, cute::_, cute::_, step));
+ static_assert(cute::size(src) == 4);
+ CUTLASS_PRAGMA_UNROLL
+ for (int i = 0; i < cute::size(src); i += 2) {
+ float4 va = *reinterpret_cast<const float4*>(&src(i)), vb = *reinterpret_cast<const float4*>(&src(i + 1));
+ float norm = 1.0f;
+ if constexpr (rms) {
+ // Preserve the standalone RMSNorm reduction tree across four 32-element groups.
+ float sum = __fadd_rn(sum32(va), sum32(vb));
+ sum = __fadd_rn(sum, __shfl_xor_sync(0xffffffff, sum, 8));
+ norm = rsqrtf(sum / 128.f + extra->eps);
+ }
+ CUTLASS_PRAGMA_UNROLL
+ for (int j = 0; j < 2; ++j) {
+ auto coord = coords(i + j);
+ int row = cute::get<0>(coord), col = cute::get<1>(coord);
+ int heads = cute::get<1>(this->problem_shape) / 128;
+ if (cute::elem_less(coord, this->problem_shape)) {
+ float4 value = j ? vb : va;
+ int local = col % 128;
+ int64_t token = row % extra->tokens, batch = row / extra->tokens, head = col / 128;
+ float4 r = value;
+ if constexpr (rms) {
+ float4 weight = ((const float4*)extra->weight)[local / 4];
+ value.x = __fmul_rn(__fmul_rn(norm, value.x), weight.x);
+ value.y = __fmul_rn(__fmul_rn(norm, value.y), weight.y);
+ value.z = __fmul_rn(__fmul_rn(norm, value.z), weight.z);
+ value.w = __fmul_rn(__fmul_rn(norm, value.w), weight.w);
+ float4 p0 = ((const float4*)extra->pe)[token * 64 + local / 2], p1 = ((const float4*)extra->pe)[token * 64 + local / 2 + 1];
+ r.x = __fadd_rn(__fmul_rn(value.x, p0.x), __fmul_rn(value.y, p0.y));
+ r.y = __fadd_rn(__fmul_rn(value.x, p0.z), __fmul_rn(value.y, p0.w));
+ r.z = __fadd_rn(__fmul_rn(value.z, p1.x), __fmul_rn(value.w, p1.y));
+ r.w = __fadd_rn(__fmul_rn(value.z, p1.z), __fmul_rn(value.w, p1.w));
+ }
+ int64_t offset = ((batch * heads + head) * (extra->tokens + extra->prefix_tokens) + token + extra->prefix_tokens) * 128 + local;
+ if constexpr (std::is_same<Out, half>::value) {
+ half2 a = __floats2half2_rn(__fmul_rn(r.x, extra->scale), __fmul_rn(r.y, extra->scale));
+ half2 b = __floats2half2_rn(__fmul_rn(r.z, extra->scale), __fmul_rn(r.w, extra->scale));
+ half2* dst = (half2*)(extra->half_output + offset);
+ dst[0] = a;
+ dst[1] = b;
+ } else {
+ *((float4*)(this->params_ptr->ptr_aux + offset)) = r;
+ }
+ }
+ }
+ }
+ }
+ };
+ template <class Shape>
+ CUTLASS_DEVICE auto get_callbacks(cutlass::gemm::GemmCoord tile, int tid, Shape shape) {
+ auto cb = Base::get_callbacks(tile, tid, shape);
+ return Callbacks<decltype(cb)>(cb, extra);
+ }
+};
+
+template <class Out, bool RMS = true>
+struct ProjectionKernel {
+ using Tile = cutlass::gemm::GemmShape<128, 256, 64>;
+ using Warp = cutlass::gemm::GemmShape<64, 64, 64>;
+ using Swizzle = cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>;
+ using Map = cutlass::epilogue::threadblock::OutputTileThreadLayout<Tile, Warp, float, 4, 1>;
+ using Col = cutlass::epilogue::threadblock::VisitorColBroadcast<Map, float, cute::Stride<cute::_1, cute::_0, int>>;
+ using Row = cutlass::epilogue::threadblock::VisitorRowBroadcast<Map, float, cute::Stride<cute::_0, cute::_1, int>>;
+ using Mul = cutlass::epilogue::threadblock::VisitorCompute<cutlass::multiplies, float, float, cutlass::FloatRoundStyle::round_to_nearest>;
+ using ScaleX = cutlass::epilogue::threadblock::Sm80EVT<Mul, cutlass::epilogue::threadblock::VisitorAccFetch, Col>;
+ using ScaleW = cutlass::epilogue::threadblock::Sm80EVT<Mul, ScaleX, Row>;
+ using Store = I8RopeStore<Map, Out, RMS>;
+ using EVT = cutlass::epilogue::threadblock::Sm80EVT<Store, ScaleW>;
+ using Kernel = typename cutlass::gemm::kernel::DefaultGemmWithVisitor<
+ int8_t,
+ cutlass::layout::RowMajor,
+ cutlass::ComplexTransform::kNone,
+ 16,
+ int8_t,
+ cutlass::layout::ColumnMajor,
+ cutlass::ComplexTransform::kNone,
+ 16,
+ float,
+ cutlass::layout::RowMajor,
+ 4,
+ int,
+ float,
+ cutlass::arch::OpClassTensorOp,
+ cutlass::arch::Sm80,
+ Tile,
+ Warp,
+ cutlass::gemm::GemmShape<16, 8, 32>,
+ EVT,
+ Swizzle,
+ 4,
+ cutlass::arch::OpMultiplyAddSaturate>::GemmKernel;
+ using Gemm = cutlass::gemm::device::GemmUniversalAdapter<Kernel>;
+};
+
+static float param_f32(const ggml_tensor* tensor, int index) {
+ float result;
+ std::memcpy(&result, tensor->op_params + index, sizeof(result));
+ return result;
+}
+
+static bool device_supports(ggml_backend_dev_t device) {
+#ifdef SD_CUDA_FORCE_CUBLAS
+ (void)device;
+ return false;
+#else
+ static std::mutex mutex;
+ static std::unordered_map<ggml_backend_dev_t, bool> cache;
+ std::lock_guard<std::mutex> lock(mutex);
+ auto found = cache.find(device);
+ if (found != cache.end())
+ return found->second;
+ auto reg = ggml_backend_dev_backend_reg(device);
+ auto query = reinterpret_cast<ggml_ext_cuda_device_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_device_v1"));
+ int ordinal = query ? query(device) : -1;
+ cudaDeviceProp props = {};
+ size_t shared = std::max({sizeof(typename ProjectionKernel<float>::Kernel::SharedStorage),
+ sizeof(typename ProjectionKernel<half>::Kernel::SharedStorage), sizeof(typename ProjectionKernel<half, false>::Kernel::SharedStorage)});
+ bool supported = ordinal >= 0 && cudaGetDeviceProperties(&props, ordinal) == cudaSuccess && props.major >= 8 && props.sharedMemPerBlockOptin >= shared;
+ cache.emplace(device, supported);
+ return supported;
+#endif
+}
+
+static bool aligned(const ggml_tensor* tensor) {
+ size_t offset = 0;
+ while (tensor->view_src) {
+ offset += tensor->view_offs;
+ tensor = tensor->view_src;
+ }
+ return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % 16 == 0;
+}
+
+static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) {
+ SDI8Operation op;
+ if (!sd_ggml_i8_describe(mat, &op) || !op.packed || !op.weight_scale || op.bias)
+ return false;
+ const auto* w = op.weight;
+ const auto* x = op.input;
+ const auto* ws = op.weight_scale;
+ if (mat->type != GGML_TYPE_F32 ||
+ w->type != GGML_TYPE_I8 || x->type != GGML_TYPE_I8 || ws->type != GGML_TYPE_F32 ||
+ w->ne[2] != 1 || w->ne[3] != 1 || mat->ne[0] != w->ne[1] || ggml_nelements(ws) != mat->ne[0] ||
+ !ggml_is_contiguous(mat) || !ggml_is_contiguous(w) || !ggml_is_contiguous(x) || !ggml_is_contiguous(ws) ||
+ !aligned(w) || !aligned(x) || !aligned(ws) || !aligned(mat) ||
+ view->op != GGML_OP_RESHAPE || view->src[0] != mat || view->view_offs != 0 || !ggml_is_contiguous(view) ||
+ view->ne[0] != 128 || view->ne[0] * view->ne[1] != mat->ne[0] || ggml_nelements(view) != ggml_nelements(mat))
+ return false;
+ const int64_t rows = ggml_nrows(mat), n = mat->ne[0], k = w->ne[0];
+ if (rows < 4096 || rows > INT_MAX || n < 4096 || n > INT_MAX || k < 4096 || k > INT_MAX || n % 256 || k % 64 || n / 256 > 65535 ||
+ x->ne[0] != k || x->ne[1] != (rows + 3) / 4 * 4 + (rows * 4 + k - 1) / k || x->ne[2] != 1 || x->ne[3] != 1)
+ return false;
+ return device_supports(device);
+}
+
+bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+ if (!projection_supports(device, mat, view) || norm->src[0] != view ||
+ !sd_cuda::rope_fusion_supports(norm, mul, rope) || rope->src[1]->ne[2] != 64 || !aligned(mul->src[1]) || !aligned(rope->src[1]))
+ return false;
+ int32_t interleaved;
+ std::memcpy(&interleaved, ggml_ext_op_attributes(rope), sizeof(interleaved));
+ return interleaved != 0;
+}
+
+template <class Out, bool RMS = true>
+static cudaError_t projection_launch(cudaStream_t stream, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope, void* output, float scale, int64_t prefix_tokens) {
+ using P = ProjectionKernel<Out, RMS>;
+ using Gemm = typename P::Gemm;
+ const int rows = int(ggml_nrows(mat)), n = int(mat->ne[0]), k = int(mat->src[0]->ne[0]);
+ const auto* x = static_cast<const int8_t*>(mat->src[1]->data);
+ const auto* w = static_cast<const int8_t*>(mat->src[0]->data);
+ const auto* xs = reinterpret_cast<const float*>(x + int64_t(k) * ((int64_t(rows) + 3) / 4 * 4));
+ const auto* ws = static_cast<const float*>(mat->src[2]->data);
+ const auto* nw = RMS ? static_cast<const float*>(mul->src[1]->data) : nullptr;
+ const auto* pe = RMS ? static_cast<const float*>(rope->src[1]->data) : nullptr;
+ // The visitor forms FP32 addresses even when its only writes target the FP16 output.
+ float* address_space = std::is_same<Out, half>::value ? static_cast<float*>(mat->data) : static_cast<float*>(output);
+ if ((reinterpret_cast<uintptr_t>(x) | reinterpret_cast<uintptr_t>(w) | reinterpret_cast<uintptr_t>(xs) |
+ reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(nw) | reinterpret_cast<uintptr_t>(pe) |
+ reinterpret_cast<uintptr_t>(address_space)) %
+ 16 ||
+ reinterpret_cast<uintptr_t>(output) % 8)
+ return cudaErrorInvalidValue;
+ typename P::EVT::Arguments evt = {{{{}, {xs, 0.f, {}}, {}}, {ws, 0.f, {}}, {}},
+ {{address_space, {n, cute::_1{}, int64_t(rows) * n}}, nw, pe, RMS ? param_f32(norm, 0) : 0.f, scale, view->ne[2], static_cast<half*>(output), prefix_tokens}};
+ typename Gemm::Arguments args(cutlass::gemm::GemmUniversalMode::kGemm, {rows, n, k}, 1, evt, x, w, nullptr, nullptr,
+ int64_t(rows) * k, int64_t(n) * k, 0, 0, k, k, 0, 0);
+ Gemm gemm;
+ if (gemm.can_implement(args) != cutlass::Status::kSuccess)
+ return cudaErrorInvalidValue;
+ if (gemm.initialize(args, nullptr, stream) != cutlass::Status::kSuccess || gemm(stream) != cutlass::Status::kSuccess)
+ return cudaErrorUnknown;
+ return cudaGetLastError();
+}
+
+cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale) {
+ return projection_launch<half>(stream, nodes[0], nodes[1], nodes[2], nodes[3], nodes[4], output, scale, 0);
+}
+
+static const ggml_tensor* node_at(const ggml_cgraph* graph, int index) {
+ return ggml_graph_node(const_cast<ggml_cgraph*>(graph), index);
+}
+
+static int concat_match(const ggml_cgraph* graph, int first, bool key) {
+ const int offset = key ? 5 : 2, count = key ? 7 : 5;
+ if (first + count > ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
+ return 0;
+ const auto* concat = node_at(graph, first + offset);
+ const auto* copy = node_at(graph, first + count - 1);
+ if (concat->op != GGML_OP_CONCAT || concat->op_params[0] != (key ? 1 : 2) || concat->src[1] != node_at(graph, first + offset - 1) ||
+ copy->op != GGML_OP_CPY || copy->src[0] != node_at(graph, first + count - 2) || copy->type != GGML_TYPE_F16 ||
+ !ggml_is_contiguous(copy) || !ggml_are_same_shape(copy, copy->src[0]))
+ return 0;
+ for (int i = 0; i < 2; ++i)
+ if (!concat->src[i] || concat->src[i]->type != GGML_TYPE_F32 || !ggml_is_contiguous(concat->src[i]))
+ return 0;
+ const auto* view = node_at(graph, first + 1);
+ const auto* prefix = concat->src[0];
+ if (prefix->ne[key ? 1 : 2] > INT_MAX - view->ne[2] || (ggml_nelements(prefix) + 255) / 256 > INT_MAX)
+ return 0;
+ if (!key) {
+ const auto* permute = node_at(graph, first + 3);
+ if (permute->op != GGML_OP_PERMUTE || permute->src[0] != concat || permute->op_params[0] != 0 || permute->op_params[1] != 2 ||
+ permute->op_params[2] != 1 || permute->op_params[3] != 3)
+ return 0;
+ }
+ return count;
+}
+
+static int projection_match(ggml_backend_dev_t device, const ggml_cgraph* graph, int first, void*) {
+ static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION");return value&&std::atoi(value)!=0; }();
+ if (disabled || first + 4 >= ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph)))
+ return 0;
+ const auto* mat = node_at(graph, first);
+ const auto* view = node_at(graph, first + 1);
+ const auto* norm = node_at(graph, first + 2);
+ const bool key = norm->op == GGML_OP_RMS_NORM;
+ if (key) {
+ if (!sd_cuda_i8_rope_supports(device, mat, view, norm, node_at(graph, first + 3), node_at(graph, first + 4)))
+ return 0;
+ if (int count = concat_match(graph, first, true))
+ return count;
+ if (first + 6 < ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph))) {
+ const auto* rope = node_at(graph, first + 4);
+ const auto* scale = node_at(graph, first + 5);
+ const auto* copy = node_at(graph, first + 6);
+ if (scale->op == GGML_OP_SCALE && scale->src[0] == rope && param_f32(scale, 1) == 0.f && copy->op == GGML_OP_CPY && copy->src[0] == scale &&
+ copy->type == GGML_TYPE_F16 && ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+ return 7;
+ }
+ return 5;
+ }
+ return projection_supports(device, mat, view) ? concat_match(graph, first, false) : 0;
+}
+
+template <bool Aligned>
+static __global__ void prefix_copy(const float* input, half2* output, int64_t count4, int64_t prefix, int64_t tokens, int64_t heads, bool permuted) {
+ const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+ if (i >= count4)
+ return;
+ const int64_t col = i % 32, row = i / 32;
+ const int64_t token = row % prefix, head = (row / prefix) % heads, batch = row / (prefix * heads);
+ const int64_t source_row = permuted ? (batch * prefix + token) * heads + head : (batch * heads + head) * prefix + token;
+ const float* source = input + source_row * 128 + col * 4;
+ float4 value;
+ if constexpr (Aligned)
+ value = *reinterpret_cast<const float4*>(source);
+ else
+ value = make_float4(source[0], source[1], source[2], source[3]);
+ const int64_t target = ((batch * heads + head) * (prefix + tokens) + token) * 32 + col;
+ output[2 * target] = __floats2half2_rn(value.x, value.y);
+ output[2 * target + 1] = __floats2half2_rn(value.z, value.w);
+}
+
+static ggml_status projection_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* workspace, void*) {
+ *workspace = ggml_nbytes(nodes[count - 1]);
+ return GGML_STATUS_SUCCESS;
+}
+static void projection_destroy(void*, void*) {}
+
+static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+}
+
+static ggml_status projection_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+ const auto* mat = nodes[0];
+ const auto* view = nodes[1];
+ const bool key = nodes[2]->op == GGML_OP_RMS_NORM;
+ const auto* concat = nodes[key ? count == 7 ? 5 : 4 : 2];
+ const bool joined = concat->op == GGML_OP_CONCAT;
+ auto* dst = nodes[count - 1];
+ auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+ auto query = reinterpret_cast<ggml_ext_cuda_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_cuda_context_v1"));
+ ggml_ext_cuda_context context;
+ if (!query || !query(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<cudaStream_t>(context.stream);
+ bool alias = reinterpret_cast<uintptr_t>(dst->data) % (dst->type == GGML_TYPE_F32 ? 16 : 8) != 0;
+ for (const auto* input : {mat->src[0], mat->src[1], mat->src[2], key ? nodes[3]->src[1] : nullptr, key ? nodes[4]->src[1] : nullptr, joined ? concat->src[0] : nullptr})
+ if (input)
+ alias |= overlaps(dst, input);
+ void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+ int64_t prefix = joined ? concat->src[0]->ne[key ? 1 : 2] : 0;
+ float scale = key && count == 7 && !joined ? param_f32(nodes[5], 0) : 1.f;
+ cudaError_t error;
+ if (dst->type == GGML_TYPE_F32)
+ error = projection_launch<float>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, 0);
+ else if (key)
+ error = projection_launch<half>(stream, mat, view, nodes[2], nodes[3], nodes[4], output, scale, prefix);
+ else
+ error = projection_launch<half, false>(stream, mat, view, nullptr, nullptr, nullptr, output, scale, prefix);
+ if (error == cudaSuccess && joined && prefix > 0) {
+ const auto* input = concat->src[0];
+ const int64_t count4 = ggml_nelements(input) / 4;
+ auto kernel = reinterpret_cast<uintptr_t>(input->data) % 16 == 0 ? prefix_copy<true> : prefix_copy<false>;
+ kernel<<<static_cast<unsigned int>((count4 + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(input->data), static_cast<half2*>(output), count4, prefix, view->ne[2], view->ne[1], !key);
+ error = cudaGetLastError();
+ }
+ if (error == cudaSuccess && alias)
+ error = cudaMemcpyAsync(dst->data, output, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, stream);
+ return error == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+}
+
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.i8_projection.cuda.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_i8_op();
+ desc.priority = 350;
+ desc.match = projection_match;
+ desc.prepare = projection_prepare;
+ desc.execute = projection_execute;
+ desc.destroy = projection_destroy;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+}
diff --git a/src/core/cuda/i8_projection.h b/src/core/cuda/i8_projection.h
new file mode 100644
index 0000000..fb815d0
--- /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 <cuda_runtime.h>
+#include "ggml-ext.h"
+
+bool sd_cuda_i8_rope_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view, const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope);
+cudaError_t sd_cuda_i8_rope_launch(cudaStream_t stream, ggml_tensor* const* nodes, void* output, float scale);
+void sd_ggml_register_cuda_i8_projection(ggml_backend_dev_t device);
+
+#endif
diff --git a/src/core/cuda/int8.cu b/src/core/cuda/int8.cu
new file mode 100644
index 0000000..45f6ecd
--- /dev/null
+++ b/src/core/cuda/int8.cu
@@ -0,0 +1,739 @@
+#include <cstdlib>
+#include "convrot_kernel.h"
+#include "core/ggml_extensions.h"
+#ifdef SD_CUDA_I8_CUTLASS
+#include "core/cuda/i8_gemm.h"
+#else
+struct SDI8CudaGemmPlan;
+#endif
+
+#include "runtime.h"
+#if defined(SD_CUDA_VENDOR_HIP)
+#include <hipblas/hipblas.h>
+#elif !defined(SD_CUDA_VENDOR_MUSA)
+#include <cublas_v2.h>
+#endif
+#include <algorithm>
+#include <climits>
+#include <mutex>
+#include <unordered_map>
+
+namespace SD_CUDA_NAMESPACE {
+ namespace sd_i8_cuda {
+ template <int Width = 32>
+ static __device__ __forceinline__ float warp_reduce_max(float value) {
+#pragma unroll
+ for (int offset = Width / 2; offset > 0; offset >>= 1)
+ value = fmaxf(value, cuda_shuffle_xor(value, offset, Width));
+ return value;
+ }
+ static __global__ void quantize_rowwise_i8_cuda(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows) {
+ const int64_t row = blockIdx.x;
+ const int tid = threadIdx.x;
+ if (row >= rows) {
+ return;
+ }
+
+ __shared__ float maxima[256];
+ const float* src_row = src + row * k;
+ int8_t* dst_row = dst + row * k;
+ float local_max = 0.0f;
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ local_max = fmaxf(local_max, fabsf(src_row[i]));
+ }
+ maxima[tid] = local_max;
+ __syncthreads();
+
+ for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+ if (tid < stride) {
+ maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+ }
+ __syncthreads();
+ }
+
+ const float scale = maxima[0] / 127.0f;
+ if (tid == 0) {
+ scales[row] = scale;
+ }
+ if (scale == 0.0f) {
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ dst_row[i] = 0;
+ }
+ return;
+ }
+
+ const float inv_scale = 1.0f / scale;
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ int value = __float2int_rn(src_row[i] * inv_scale);
+ value = max(-127, min(127, value));
+ dst_row[i] = (int8_t)value;
+ }
+ }
+ static __device__ __forceinline__ float convrot_h4_cuda(float a, float b, float c, float d, int index) {
+ a = index == 3 ? -a : a;
+ b = index == 2 ? -b : b;
+ c = index == 1 ? -c : c;
+ d = index == 0 ? -d : d;
+ return __fadd_rn(__fadd_rn(__fadd_rn(a, b), c), d);
+ }
+ static __global__ void quantize_convrot(const float* src, int8_t* dst, float* scales, int64_t k, int64_t rows, int64_t dst_stride) {
+ constexpr int group_size = 256;
+ const int64_t row = blockIdx.x;
+ const int tid = threadIdx.x;
+ if (row >= rows)
+ return;
+ extern __shared__ float values[];
+ __shared__ float maxima[32];
+ const float* src_row = src + row * k;
+ int8_t* dst_row = dst + row * dst_stride;
+ {
+ const int lane = tid % 32;
+ for (int64_t group = tid / 32; group < k / group_size; group += blockDim.x / 32) {
+ float v[8];
+#pragma unroll
+ for (int j = 0; j < 8; ++j) {
+ const int64_t i = group * group_size + j * 32 + lane;
+ const float value = src_row[i];
+ v[j] = value * (1.0f / 16.0f);
+ }
+ // One logical warp owns H256, keeping every butterfly in registers.
+#pragma unroll
+ for (int stride = 1; stride <= 4; stride *= 4) {
+ const int base = lane & ~(3 * stride);
+#pragma unroll
+ for (int j = 0; j < 8; ++j) {
+ const float a = cuda_shuffle(v[j], base, 32);
+ const float b = cuda_shuffle(v[j], base + stride, 32);
+ const float c = cuda_shuffle(v[j], base + 2 * stride, 32);
+ const float d = cuda_shuffle(v[j], base + 3 * stride, 32);
+ v[j] = convrot_h4_cuda(a, b, c, d, (lane / stride) % 4);
+ }
+ }
+#pragma unroll
+ for (int j = 0; j < 8; j += 2) {
+ const float a = cuda_shuffle(v[j], lane % 16, 32);
+ const float b = cuda_shuffle(v[j], lane % 16 + 16, 32);
+ const float c = cuda_shuffle(v[j + 1], lane % 16, 32);
+ const float d = cuda_shuffle(v[j + 1], lane % 16 + 16, 32);
+ v[j] = convrot_h4_cuda(a, b, c, d, lane / 16);
+ v[j + 1] = convrot_h4_cuda(a, b, c, d, 2 + lane / 16);
+ }
+#pragma unroll
+ for (int j = 0; j < 2; ++j) {
+ const float a = v[j], b = v[j + 2], c = v[j + 4], d = v[j + 6];
+#pragma unroll
+ for (int t = 0; t < 4; ++t) {
+ values[group * group_size + (j + 2 * t) * 32 + lane] = convrot_h4_cuda(a, b, c, d, t);
+ }
+ }
+ }
+ __syncthreads();
+ }
+ float local_max = 0.0f;
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ local_max = fmaxf(local_max, fabsf(values[i]));
+ }
+ local_max = warp_reduce_max<32>(local_max);
+ const int lane = tid & 31;
+ const int warp = tid >> 5;
+ if (lane == 0) {
+ maxima[warp] = local_max;
+ }
+ __syncthreads();
+
+ if (warp == 0) {
+ local_max = lane < blockDim.x / 32 ? maxima[lane] : 0.0f;
+ local_max = warp_reduce_max<32>(local_max);
+ if (lane == 0) {
+ maxima[0] = local_max / 127.0f;
+ if (scales != nullptr) {
+ scales[row] = maxima[0];
+ } else {
+ *((float*)(dst_row + k)) = maxima[0];
+ }
+ }
+ }
+ __syncthreads();
+
+ const float row_scale = maxima[0];
+ const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+ for (int64_t i = tid; i < k; i += blockDim.x) {
+ int value = row_scale == 0.0f ? 0 : __float2int_rn(values[i] * inv_scale);
+ value = max(-127, min(127, value));
+ dst_row[i] = (int8_t)value;
+ }
+ }
+
+ static __global__ void convrot_group_amax_i8_cuda(
+ const float* src,
+ float* partial_max,
+ int64_t k,
+ int64_t groups,
+ int64_t total_groups) {
+ const int64_t index = blockIdx.x;
+ const int tid = threadIdx.x;
+ if (index >= total_groups) {
+ return;
+ }
+
+ __shared__ float values[256];
+ __shared__ float maxima[2];
+ const int64_t row = index / groups;
+ const int64_t group = index - row * groups;
+ const float* input = src + row * k + group * 256;
+ values[tid] = input[tid] * (1.0f / 16.0f);
+ values[tid + 64] = input[tid + 64] * (1.0f / 16.0f);
+ values[tid + 128] = input[tid + 128] * (1.0f / 16.0f);
+ values[tid + 192] = input[tid + 192] * (1.0f / 16.0f);
+ __syncthreads();
+
+#pragma unroll
+ for (int stride = 1; stride < 256; stride *= 4) {
+ const int base = (tid / stride) * 4 * stride + tid % stride;
+ const int i0 = base;
+ const int i1 = i0 + stride;
+ const int i2 = i1 + stride;
+ const int i3 = i2 + stride;
+ const float a = values[i0];
+ const float b = values[i1];
+ const float c = values[i2];
+ const float d = values[i3];
+ values[i0] = a + b + c - d;
+ values[i1] = a + b - c + d;
+ values[i2] = a - b + c + d;
+ values[i3] = -a + b + c + d;
+ __syncthreads();
+ }
+
+ float value = fmaxf(fabsf(values[tid]), fabsf(values[tid + 64]));
+ value = fmaxf(value, fabsf(values[tid + 128]));
+ value = fmaxf(value, fabsf(values[tid + 192]));
+ value = warp_reduce_max<32>(value);
+ if ((tid & 31) == 0) {
+ maxima[tid >> 5] = value;
+ }
+ __syncthreads();
+ if (tid == 0) {
+ partial_max[index] = fmaxf(maxima[0], maxima[1]);
+ }
+ }
+
+ static __global__ void reduce_convrot_row_amax_i8_cuda(
+ const float* partial_max,
+ float* scales,
+ int64_t groups,
+ int64_t rows) {
+ const int64_t row = blockIdx.x;
+ const int tid = threadIdx.x;
+ if (row >= rows) {
+ return;
+ }
+
+ __shared__ float maxima[256];
+ float value = 0.0f;
+ for (int64_t group = tid; group < groups; group += blockDim.x) {
+ value = fmaxf(value, partial_max[row * groups + group]);
+ }
+ maxima[tid] = value;
+ __syncthreads();
+ for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
+ if (tid < stride) {
+ maxima[tid] = fmaxf(maxima[tid], maxima[tid + stride]);
+ }
+ __syncthreads();
+ }
+ if (tid == 0) {
+ scales[row] = maxima[0] / 127.0f;
+ }
+ }
+
+ static __global__ void convrot_group_quantize_i8_cuda(
+ const float* src,
+ int8_t* dst,
+ const float* scales,
+ int64_t k,
+ int64_t groups,
+ int64_t total_groups,
+ int64_t dst_stride) {
+ const int64_t index = blockIdx.x;
+ const int tid = threadIdx.x;
+ if (index >= total_groups) {
+ return;
+ }
+
+ __shared__ float values[256];
+ const int64_t row = index / groups;
+ const int64_t group = index - row * groups;
+ const float* input = src + row * k + group * 256;
+ int8_t* output = dst + row * dst_stride + group * 256;
+ values[tid] = input[tid] * (1.0f / 16.0f);
+ values[tid + 64] = input[tid + 64] * (1.0f / 16.0f);
+ values[tid + 128] = input[tid + 128] * (1.0f / 16.0f);
+ values[tid + 192] = input[tid + 192] * (1.0f / 16.0f);
+ __syncthreads();
+
+#pragma unroll
+ for (int stride = 1; stride < 256; stride *= 4) {
+ const int base = (tid / stride) * 4 * stride + tid % stride;
+ const int i0 = base;
+ const int i1 = i0 + stride;
+ const int i2 = i1 + stride;
+ const int i3 = i2 + stride;
+ const float a = values[i0];
+ const float b = values[i1];
+ const float c = values[i2];
+ const float d = values[i3];
+ values[i0] = a + b + c - d;
+ values[i1] = a + b - c + d;
+ values[i2] = a - b + c + d;
+ values[i3] = -a + b + c + d;
+ __syncthreads();
+ }
+
+ const float row_scale = scales[row];
+ const float inv_scale = row_scale == 0.0f ? 0.0f : 1.0f / row_scale;
+ const int indices[4] = {tid, tid + 64, tid + 128, tid + 192};
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ int value = row_scale == 0.0f ? 0 : __float2int_rn(values[indices[i]] * inv_scale);
+ value = max(-127, min(127, value));
+ output[indices[i]] = (int8_t)value;
+ }
+ if (group == 0 && tid == 0 && dst_stride != k) {
+ *((float*)(dst + row * dst_stride + k)) = row_scale;
+ }
+ }
+
+ static __global__ void dequantize_i32_rows_cuda(
+ const int32_t* src,
+ const float* scales,
+ const float* weight_scales,
+ const float* bias,
+ float* dst,
+ int64_t n,
+ int64_t rows) {
+ const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+ const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535;
+ if (output < n && row < rows) {
+ const int64_t index = row * n + output;
+ const float activation_scale = scales[row];
+ float value = (float)src[index] * activation_scale;
+ if (weight_scales != nullptr) {
+ value *= weight_scales[output];
+ }
+ if (bias != nullptr) {
+ value += bias[output];
+ }
+ dst[index] = value;
+ }
+ }
+
+ static __global__ void dequantize_i32_rows_vec4_cuda(
+ const int4* src,
+ const float* scales,
+ const float4* weight_scales,
+ const float4* bias,
+ float4* dst,
+ int64_t n4,
+ int64_t rows) {
+ const int64_t output = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+ const int64_t row = blockIdx.y + int64_t(blockIdx.z) * 65535;
+ if (output >= n4 || row >= rows) {
+ return;
+ }
+ const int64_t index = row * n4 + output;
+ const int4 v = src[index];
+ const float scale = scales[row];
+ float4 value = make_float4(v.x * scale, v.y * scale, v.z * scale, v.w * scale);
+ if (weight_scales != nullptr) {
+ const float4 w = weight_scales[output];
+ value.x *= w.x;
+ value.y *= w.y;
+ value.z *= w.z;
+ value.w *= w.w;
+ }
+ if (bias != nullptr) {
+ const float4 b = bias[output];
+ value.x += b.x;
+ value.y += b.y;
+ value.z += b.z;
+ value.w += b.w;
+ }
+ dst[index] = value;
+ }
+
+ static bool aligned(const ggml_tensor* tensor, size_t alignment = 4) {
+ if (!tensor)
+ return true;
+ size_t offset = 0;
+ while (tensor->view_src) {
+ offset += tensor->view_offs;
+ tensor = tensor->view_src;
+ }
+ return (reinterpret_cast<uintptr_t>(tensor->data) + offset) % alignment == 0;
+ }
+ static bool supports(ggml_backend_dev_t device, const ggml_tensor* node, void*) {
+ SDI8Operation op;
+ if (!sd_ggml_i8_describe(node, &op) || !aligned(node))
+ return false;
+ const int64_t k = op.weight->ne[0], n = op.weight->ne[1], rows = ggml_nrows(node);
+ if (k > INT_MAX || n > INT_MAX || rows > INT_MAX - 3 || k % 4 || n % 4 || (op.input->type == GGML_TYPE_F32 && op.group_size && (op.group_size != 256 || rows > INT_MAX / (k / 256))))
+ return false;
+ for (auto* input : node->src)
+ if (!aligned(input))
+ return false;
+ const auto info = cuda_device_info(device);
+ return info.int8_blas &&
+ ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_blas_v1");
+ }
+ 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<ggml_cgraph*>(graph);
+ if (fusion_disabled() || first + 1 >= ggml_graph_n_nodes(g))
+ return 0;
+ const auto* mat = ggml_graph_node(g, first);
+ SDI8Operation op;
+ if (!supports(device, mat, nullptr) || !sd_ggml_i8_describe(mat, &op) || !op.weight_scale || op.bias)
+ return 0;
+ const auto* next = ggml_graph_node(g, first + 1);
+ if (next->op == GGML_OP_CONCAT && next->src[1] == mat && next->src[0] != mat && next->op_params[0] == 1 &&
+ next->type == GGML_TYPE_F32 && mat->ne[2] == 1 && mat->ne[3] == 1 &&
+ next->ne[2] == 1 && next->ne[3] == 1 && ggml_is_contiguous(next) && ggml_is_contiguous(next->src[0]))
+ return 2;
+ if (next->op != GGML_OP_GLU || ggml_get_glu_op(next) != GGML_GLU_OP_SWIGLU || next->op_params[1] != 0 ||
+ next->src[0] != mat || next->src[1] || next->type != GGML_TYPE_F32 || !ggml_is_contiguous(next) ||
+ ggml_nelements(next) > int64_t(INT_MAX) * 256)
+ return 0;
+ if (first + 2 < ggml_graph_n_nodes(g)) {
+ const auto* quant = ggml_graph_node(g, first + 2);
+ 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<Plan*>(state);
+#ifdef SD_CUDA_I8_CUTLASS
+ if (p)
+ sd_cuda_i8_gemm_free(p->gemm);
+#endif
+ delete p;
+ }
+ static bool overlaps(const ggml_tensor* a, const ggml_tensor* b) {
+ if (!b)
+ return false;
+ const auto x = reinterpret_cast<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+ }
+ 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<int, bool> configured;
+ {
+ std::lock_guard<std::mutex> lock(mutex);
+ if (!configured[p.device.device]) {
+ auto error = cuda_set_shared_memory(quantize_convrot, int(p.device.shared_memory - 128));
+ if (error != cuda_success)
+ return error;
+ configured[p.device.device] = true;
+ }
+ }
+ int threads = p.k <= 1024 ? 128 : p.k <= 8192 ? 256
+ : 1024;
+ quantize_convrot<<<unsigned(p.rows), threads, shared, stream>>>(src, dst, scales, p.k, p.rows, p.k);
+ } else {
+ const int64_t groups = p.k / 256, total_groups = p.rows * groups;
+ convrot_group_amax_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, partial, p.k, groups, total_groups);
+ reduce_convrot_row_amax_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(partial, scales, groups, p.rows);
+ convrot_group_quantize_i8_cuda<<<unsigned(total_groups), 64, 0, stream>>>(src, dst, scales, p.k, groups, total_groups, p.k);
+ }
+ return cuda_last_error();
+ }
+ template <bool Dequantize>
+ static __global__ void swiglu(const float* src, const float* xs, const float* ws, float* dst, int64_t n, int64_t rows) {
+ const int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x;
+ const int64_t half = n / 2;
+ if (i >= rows * half)
+ return;
+ const int64_t row = i / half, col = i % half;
+ float gate, up;
+ if constexpr (Dequantize) {
+ gate = float(reinterpret_cast<const int32_t*>(src)[row * n + col]) * xs[row];
+ up = float(reinterpret_cast<const int32_t*>(src)[row * n + half + col]) * xs[row];
+ gate *= ws[col];
+ up *= ws[half + col];
+ } else {
+ gate = src[row * n + col];
+ up = src[row * n + half + col];
+ }
+ dst[i] = __fmul_rn(gate / (1.f + expf(-gate)), up);
+ }
+ template <bool Swiglu, bool Dequantize, bool Rotated>
+ static CUDAError quantize_epilogue(const Plan& p, const float* src, const float* xs, const float* ws, int8_t* output, CUDAStream stream) {
+ const int64_t hidden = p.n / 2;
+ auto kernel = sd_convrot::quantize<Swiglu, Dequantize, false, false, Rotated>;
+ static std::mutex mutex;
+ static std::unordered_map<int, bool> initialized;
+ {
+ std::lock_guard<std::mutex> lock(mutex);
+ if (!initialized[p.device.device]) {
+ auto error = cuda_set_shared_memory(kernel, int(p.device.shared_memory - 128));
+ if (error != cuda_success)
+ return error;
+ initialized[p.device.device] = true;
+ }
+ }
+ if (p.padded > p.rows) {
+ auto error = cuda_memset_async(output + hidden * p.rows, 0, size_t(hidden) * (p.padded - p.rows), stream);
+ if (error != cuda_success)
+ return error;
+ }
+ const int threads = Rotated && !Swiglu && !Dequantize ? (hidden <= 1024 ? 128 : hidden <= 8192 ? 256
+ : 1024)
+ : 256;
+ kernel<<<unsigned(p.rows), threads, size_t(hidden) * sizeof(float), stream>>>(src, output,
+ reinterpret_cast<float*>(output + hidden * p.padded), hidden, p.rows, hidden,
+ Swiglu ? src + hidden : nullptr, Swiglu ? p.n : hidden, Swiglu ? p.n : hidden, xs, ws, nullptr, 0.f, 1.f, 0.f, nullptr, nullptr);
+ return cuda_last_error();
+ }
+ static ggml_status execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void* state, void*) {
+ const auto& p = *static_cast<const Plan*>(state);
+ auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(execution->backend));
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context) || context.device != p.device.device)
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<CUDAStream>(context.stream);
+ auto* scratch = static_cast<char*>(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<int8_t*>(output);
+ auto* scales = reinterpret_cast<float*>(quant + p.k * p.padded);
+ error = convrot_launch(p, stream, static_cast<const float*>(dst->src[0]->data), quant, scales, reinterpret_cast<float*>(scratch + p.partial));
+ } else {
+ auto* quant = reinterpret_cast<int8_t*>(scratch + p.quant);
+ const float* scales;
+ if (op.input->type == GGML_TYPE_F32) {
+ auto* temporary_scales = reinterpret_cast<float*>(scratch + p.scales);
+ scales = temporary_scales;
+ if (op.group_size)
+ error = convrot_launch(p, stream, static_cast<const float*>(op.input->data), quant, temporary_scales, reinterpret_cast<float*>(scratch + p.partial));
+ else {
+ if (p.padded > p.rows)
+ error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+ if (error == cuda_success)
+ quantize_rowwise_i8_cuda<<<unsigned(p.rows), 256, 0, stream>>>(static_cast<const float*>(op.input->data), quant, temporary_scales, p.k, p.rows);
+ }
+ } else {
+ auto* data = static_cast<int8_t*>(op.input->data);
+ scales = op.packed ? reinterpret_cast<const float*>(data + p.k * p.padded) : static_cast<const float*>(op.input_scale->data);
+ if (op.packed || p.rows == p.padded)
+ quant = data;
+ else {
+ error = cuda_copy_async(quant, data, size_t(p.k) * p.rows, stream);
+ if (error == cuda_success)
+ error = cuda_memset_async(quant + p.k * p.rows, 0, size_t(p.k) * (p.padded - p.rows), stream);
+ }
+ }
+ if (error != cuda_success || cuda_last_error() != cuda_success)
+ return GGML_STATUS_FAILED;
+ const auto* ws = op.weight_scale ? static_cast<const float*>(op.weight_scale->data) : nullptr;
+ const auto* bias = op.bias ? static_cast<const float*>(op.bias->data) : nullptr;
+ const size_t prefix = p.kind == CONCAT ? ggml_nbytes(dst->src[0]) : 0;
+ auto* gemm_output = reinterpret_cast<float*>(static_cast<char*>(output) + prefix);
+ auto* accum = reinterpret_cast<int32_t*>(scratch + p.accum);
+ const bool intermediate = p.kind == SWIGLU_QUANTIZE || (p.kind == SWIGLU && !p.fast_swiglu);
+#ifdef SD_CUDA_I8_CUTLASS
+ if (p.gemm) {
+ auto status = sd_cuda_i8_gemm_execute(p.gemm, quant, static_cast<const int8_t*>(op.weight->data), scales, ws,
+ intermediate ? reinterpret_cast<float*>(accum) : gemm_output, scratch + p.gemm_offset, stream);
+ if (status != GGML_STATUS_SUCCESS)
+ return status;
+ } else
+#endif
+ {
+ auto blas_query = reinterpret_cast<ggml_ext_gpu_blas_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_blas_v1"));
+ auto handle = blas_query ? blas_query(execution->backend) : nullptr;
+ if (!handle)
+ return GGML_STATUS_FAILED;
+#if defined(SD_CUDA_VENDOR_MUSA)
+ return GGML_STATUS_FAILED;
+#else
+ const int32_t alpha = 1, beta = 0;
+#if defined(SD_CUDA_VENDOR_HIP)
+#if HIP_VERSION >= 60500000
+ constexpr auto compute_type = HIPBLAS_COMPUTE_32I;
+#else
+ constexpr auto compute_type = HIPBLAS_R_32I;
+#endif
+ const auto status = hipblasGemmEx(static_cast<hipblasHandle_t>(handle), HIPBLAS_OP_T, HIPBLAS_OP_N,
+ int(p.n), int(p.padded), int(p.k), &alpha,
+ op.weight->data, HIPBLAS_R_8I, int(p.k), quant, HIPBLAS_R_8I, int(p.k),
+ &beta, accum, HIPBLAS_R_32I, int(p.n), compute_type, HIPBLAS_GEMM_DEFAULT);
+ if (status != HIPBLAS_STATUS_SUCCESS)
+#else
+ const auto status = cublasGemmEx(static_cast<cublasHandle_t>(handle), CUBLAS_OP_T, CUBLAS_OP_N,
+ int(p.n), int(p.padded), int(p.k), &alpha,
+ op.weight->data, CUDA_R_8I, int(p.k), quant, CUDA_R_8I, int(p.k),
+ &beta, accum, CUDA_R_32I, int(p.n), CUBLAS_COMPUTE_32I, CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+ if (status != CUBLAS_STATUS_SUCCESS)
+#endif
+ return GGML_STATUS_FAILED;
+#endif
+ }
+ const auto* values = reinterpret_cast<const float*>(accum);
+ if (p.kind == SWIGLU_QUANTIZE) {
+ if (p.fast_swiglu)
+ error = quantize_epilogue<false, false, true>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
+ else if (p.gemm)
+ error = quantize_epilogue<true, false, false>(p, values, nullptr, nullptr, static_cast<int8_t*>(output), stream);
+ else
+ error = quantize_epilogue<true, true, false>(p, values, scales, ws, static_cast<int8_t*>(output), stream);
+ } else if (p.kind == SWIGLU && !p.fast_swiglu) {
+ const unsigned blocks = unsigned((p.rows * (p.n / 2) + 255) / 256);
+ if (p.gemm)
+ swiglu<false><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+ else
+ swiglu<true><<<blocks, 256, 0, stream>>>(values, scales, ws, static_cast<float*>(output), p.n, p.rows);
+ error = cuda_last_error();
+ } else if (!p.gemm) {
+ const unsigned gy = unsigned(std::min<int64_t>(p.rows, 65535)), gz = unsigned((p.rows + 65534) / 65535);
+ if ((reinterpret_cast<uintptr_t>(ws) | reinterpret_cast<uintptr_t>(bias) | reinterpret_cast<uintptr_t>(gemm_output)) % 16 == 0) {
+ dequantize_i32_rows_vec4_cuda<<<dim3(unsigned((p.n + 1023) / 1024), gy, gz), 256, 0, stream>>>(
+ reinterpret_cast<const int4*>(accum), scales, reinterpret_cast<const float4*>(ws), reinterpret_cast<const float4*>(bias), reinterpret_cast<float4*>(gemm_output), p.n / 4, p.rows);
+ } else {
+ dequantize_i32_rows_cuda<<<dim3(unsigned((p.n + 255) / 256), gy, gz), 256, 0, stream>>>(accum, scales, ws, bias, gemm_output, p.n, p.rows);
+ }
+ error = cuda_last_error();
+ }
+ if (error == cuda_success && prefix)
+ error = cuda_copy_async(output, dst->src[0]->data, prefix, stream);
+ }
+ if (error == cuda_success && alias)
+ error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+ return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+
+ }
+
+ void register_i8(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.linear_i8." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_i8_op();
+ desc.priority = 100;
+ desc.supports = sd_i8_cuda::supports;
+ desc.prepare = sd_i8_cuda::prepare;
+ desc.execute = sd_i8_cuda::execute;
+ desc.destroy = sd_i8_cuda::destroy;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1";
+ desc.priority = 240;
+ desc.supports = nullptr;
+ desc.match = sd_i8_cuda::match;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.priority = 100;
+ desc.match = nullptr;
+ desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
+ desc.custom_op = sd_ggml_convrot_op();
+ desc.supports = sd_i8_cuda::convrot_supports;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ }
+
+}
diff --git a/src/core/cuda/registration.cu b/src/core/cuda/registration.cu
new file mode 100644
index 0000000..d03b114
--- /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 0000000..f98c36f
--- /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 0000000..1b724b0
--- /dev/null
+++ b/src/core/cuda/rope.cu
@@ -0,0 +1,337 @@
+#include "core/cuda/rope.h"
+#include "core/ggml_extensions.h"
+
+#include <climits>
+#include <cstdlib>
+#include <cstring>
+#include <type_traits>
+
+namespace SD_CUDA_NAMESPACE {
+ static __device__ __forceinline__ float sd_warp_sum(float v) {
+ for (int mask = 16; mask > 0; mask >>= 1)
+ v = __fadd_rn(v, cuda_shuffle_xor(v, mask));
+ return v;
+ }
+ template <int D, bool interleaved, typename T>
+ static __global__ void rms_norm_mul_rope_apply_fused(
+ const float* x,
+ const float* weight,
+ const float* pe,
+ T* dst,
+ int64_t heads,
+ int64_t tokens,
+ int64_t rows,
+ int64_t nb1,
+ int64_t nb2,
+ int64_t nb3,
+ int rot_half,
+ float eps,
+ float output_scale) {
+ const int lane = threadIdx.x;
+ const int64_t row = int64_t(blockIdx.x) * blockDim.y + threadIdx.y;
+ if (row >= rows) {
+ return;
+ }
+ const int64_t head = row % heads;
+ const int64_t token = (row / heads) % tokens;
+ const int64_t batch = row / (heads * tokens);
+ const float* x_row = x + head * nb1 + token * nb2 + batch * nb3;
+ float v[D / 32], sums[8] = {};
+#pragma unroll
+ for (int j = 0; j < D / 32; ++j) {
+ v[j] = x_row[j * 32 + lane];
+ sums[j] = sd_warp_sum(__fmul_rn(v[j], v[j]));
+ }
+ // Preserve the 256-thread RMSNorm kernel's reduction tree, including inactive warps.
+ const float sum = __fadd_rn(
+ __fadd_rn(__fadd_rn(sums[0], sums[4]), __fadd_rn(sums[2], sums[6])),
+ __fadd_rn(__fadd_rn(sums[1], sums[5]), __fadd_rn(sums[3], sums[7])));
+ const float scale = rsqrtf(sum / float(D) + eps);
+ T* out = dst + ((batch * heads + head) * tokens + token) * D;
+#pragma unroll
+ for (int j = 0; j < D / 32; ++j) {
+ const int col = j * 32 + lane;
+ const float value = __fmul_rn(__fmul_rn(scale, v[j]), weight[col]);
+ float rotated = value;
+ if constexpr (interleaved) {
+ const float other = cuda_shuffle_xor(value, 1, 32);
+ if (col < 2 * rot_half) {
+ const float* m = pe + 4 * (token * rot_half + col / 2);
+ const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+ rotated = lane % 2 == 0
+ ? __fadd_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.y))
+ : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.w));
+ }
+ } else {
+ if (col < 2 * rot_half) {
+ const int partner = col < rot_half ? col + rot_half : col - rot_half;
+ const float other = __fmul_rn(__fmul_rn(scale, x_row[partner]), weight[partner]);
+ const float* m = pe + 4 * (token * rot_half + col % rot_half);
+ const float4 matrix = make_float4(m[0], m[1], m[2], m[3]);
+ rotated = col < rot_half
+ ? __fsub_rn(__fmul_rn(value, matrix.x), __fmul_rn(other, matrix.z))
+ : __fadd_rn(__fmul_rn(other, matrix.z), __fmul_rn(value, matrix.x));
+ }
+ }
+ if constexpr (std::is_same<T, half>::value) {
+ out[col] = __float2half_rn(__fmul_rn(rotated, output_scale));
+ } else {
+ out[col] = rotated;
+ }
+ }
+ }
+
+ struct RopeApplyStrides {
+ int64_t x[4];
+ int64_t pe[4];
+ };
+
+ template <bool interleaved>
+ static __global__ void rope_apply_f32(
+ const float* x,
+ const float* pe,
+ float* dst,
+ int64_t half_dim,
+ int64_t rot_half,
+ int64_t heads,
+ int64_t tokens,
+ int64_t pairs,
+ RopeApplyStrides strides) {
+ const int64_t index = (int64_t)blockIdx.x * blockDim.x + threadIdx.x;
+ if (index >= pairs) {
+ return;
+ }
+ const int64_t j = index % half_dim;
+ const int64_t row = index / half_dim;
+ const int64_t token = row % tokens;
+ const int64_t head = (row / tokens) % heads;
+ const int64_t batch = row / (tokens * heads);
+ const int64_t i0 = interleaved || j >= rot_half ? 2 * j : j;
+ const int64_t i1 = interleaved || j >= rot_half ? 2 * j + 1 : j + rot_half;
+ const float* x_row = x + head * strides.x[1] + token * strides.x[2] + batch * strides.x[3];
+ const float a = x_row[i0 * strides.x[0]];
+ const float b = x_row[i1 * strides.x[0]];
+ float* out = dst + row * half_dim * 2;
+ if (j >= rot_half) {
+ out[i0] = a;
+ out[i1] = b;
+ return;
+ }
+ const float* matrix = pe + j * strides.pe[2] + token * strides.pe[3];
+ const float c = matrix[0];
+ const float s = matrix[strides.pe[1]];
+ // Match the separate mul/add rounding of the unfused graph.
+ if (interleaved) {
+ out[i0] = __fadd_rn(__fmul_rn(a, c), __fmul_rn(b, matrix[strides.pe[0]]));
+ out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, matrix[strides.pe[0] + strides.pe[1]]));
+ } else {
+ out[i0] = __fsub_rn(__fmul_rn(a, c), __fmul_rn(b, s));
+ out[i1] = __fadd_rn(__fmul_rn(a, s), __fmul_rn(b, c));
+ }
+ }
+
+ static bool rope_interleaved(const ggml_tensor* rope) {
+ int32_t value;
+ std::memcpy(&value, ggml_ext_op_attributes(rope), sizeof(value));
+ return value != 0;
+ }
+
+ static float param_f32(const ggml_tensor* node, int index) {
+ float value;
+ std::memcpy(&value, node->op_params + index, sizeof(value));
+ return value;
+ }
+
+ static bool rope_is_node(const ggml_tensor* node) {
+ return node && ggml_ext_op_get(node) == sd_ggml_rope_op();
+ }
+
+ static bool rope_supports(ggml_backend_dev_t, const ggml_tensor* rope, void*) {
+ const auto* x = rope->src[0];
+ const auto* pe = rope->src[1];
+ if (!x || !pe || x->type != GGML_TYPE_F32 || pe->type != GGML_TYPE_F32 || rope->type != GGML_TYPE_F32 ||
+ x->ne[0] % 2 || pe->ne[0] != 2 || pe->ne[1] != 2 || pe->ne[2] <= 0 || pe->ne[2] > x->ne[0] / 2 ||
+ pe->ne[3] != x->ne[2] || !ggml_is_contiguous(rope) || rope->ne[0] != x->ne[0] ||
+ rope->ne[1] != x->ne[2] || rope->ne[2] != x->ne[1] * x->ne[3] || rope->ne[3] != 1 ||
+ (ggml_nelements(x) / 2 + 255) / 256 > INT_MAX) {
+ return false;
+ }
+ for (int i = 0; i < 4; ++i) {
+ if (x->nb[i] % sizeof(float) || pe->nb[i] % sizeof(float))
+ return false;
+ }
+ return true;
+ }
+
+ bool rope_fusion_supports(const ggml_tensor* norm, const ggml_tensor* mul, const ggml_tensor* rope) {
+ if (norm->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL || !rope_is_node(rope) ||
+ mul->src[0] != norm || rope->src[0] != mul || !rope_supports(nullptr, rope, nullptr))
+ return false;
+ const auto* x = norm->src[0];
+ const auto* weight = mul->src[1];
+ return x && weight && x->ne[0] >= 32 && x->ne[0] <= 256 && x->ne[0] % 32 == 0 &&
+ x->type == GGML_TYPE_F32 && weight->type == GGML_TYPE_F32 && norm->type == GGML_TYPE_F32 &&
+ ggml_are_same_shape(norm, mul) && ggml_are_same_shape(x, norm) &&
+ weight->ne[0] == x->ne[0] && ggml_nelements(weight) == x->ne[0] &&
+ x->nb[0] == sizeof(float) && x->nb[1] % sizeof(float) == 0 && x->nb[2] % sizeof(float) == 0 &&
+ x->nb[3] % sizeof(float) == 0 && ggml_is_contiguous(weight) && ggml_is_contiguous(rope->src[1]) &&
+ (ggml_nrows(x) + 3) / 4 <= INT_MAX;
+ }
+
+ template <typename T, bool Interleaved>
+ static CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, T* output, float scale) {
+ const auto* x = nodes[0]->src[0];
+ const auto* pe = nodes[2]->src[1];
+ const int64_t rows = ggml_nrows(x);
+#define SD_ROPE(D) \
+ case D: \
+ rms_norm_mul_rope_apply_fused<D, Interleaved, T><<<static_cast<unsigned int>((rows + 3) / 4), dim3(32, 4), 0, stream>>>( \
+ static_cast<const float*>(x->data), static_cast<const float*>(nodes[1]->src[1]->data), \
+ static_cast<const float*>(pe->data), output, x->ne[1], x->ne[2], rows, \
+ x->nb[1] / 4, x->nb[2] / 4, x->nb[3] / 4, int(pe->ne[2]), param_f32(nodes[0], 0), scale); \
+ break
+ switch (x->ne[0]) {
+ SD_ROPE(32);
+ SD_ROPE(64);
+ SD_ROPE(96);
+ SD_ROPE(128);
+ SD_ROPE(160);
+ SD_ROPE(192);
+ SD_ROPE(224);
+ SD_ROPE(256);
+ default:
+ return cuda_invalid_value;
+ }
+#undef SD_ROPE
+ return cuda_last_error();
+ }
+
+ CUDAError rms_rope_launch(CUDAStream stream, ggml_tensor* const* nodes, void* output, bool fp16, float scale) {
+ if (fp16) {
+ return rope_interleaved(nodes[2]) ? rms_rope_launch<half, true>(stream, nodes, static_cast<half*>(output), scale)
+ : rms_rope_launch<half, false>(stream, nodes, static_cast<half*>(output), scale);
+ }
+ return rope_interleaved(nodes[2]) ? rms_rope_launch<float, true>(stream, nodes, static_cast<float*>(output), scale)
+ : rms_rope_launch<float, false>(stream, nodes, static_cast<float*>(output), scale);
+ }
+
+ static int rope_match(ggml_backend_dev_t, const ggml_cgraph* graph, int first, void*) {
+ static const bool disabled = [] {const char* value=std::getenv("GGML_CUDA_DISABLE_FUSION"); return value && std::atoi(value)!=0; }();
+ if (disabled)
+ return 0;
+ auto* g = const_cast<ggml_cgraph*>(graph);
+ if (first + 2 >= ggml_graph_n_nodes(g) || !rope_fusion_supports(
+ ggml_graph_node(g, first), ggml_graph_node(g, first + 1), ggml_graph_node(g, first + 2)))
+ return 0;
+ if (first + 4 < ggml_graph_n_nodes(g)) {
+ const auto* rope = ggml_graph_node(g, first + 2);
+ const auto* scale = ggml_graph_node(g, first + 3);
+ const auto* copy = ggml_graph_node(g, first + 4);
+ if (scale->op == GGML_OP_SCALE && copy->op == GGML_OP_CPY && scale->src[0] == rope && copy->src[0] == scale &&
+ scale->type == GGML_TYPE_F32 && copy->type == GGML_TYPE_F16 && param_f32(scale, 1) == 0.f &&
+ ggml_are_same_shape(rope, scale) && ggml_are_same_shape(rope, copy) && ggml_is_contiguous(copy))
+ return 5;
+ }
+ 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<uintptr_t>(a->data), y = reinterpret_cast<uintptr_t>(b->data);
+ return x <= y ? y - x < ggml_nbytes(a) : x - y < ggml_nbytes(b);
+ }
+
+ static ggml_status rope_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int count, void*, void*) {
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ auto stream = static_cast<CUDAStream>(context.stream);
+ const auto* rope = nodes[count == 1 ? 0 : 2];
+ const auto* x = count == 1 ? rope->src[0] : nodes[0]->src[0];
+ auto* dst = nodes[count - 1];
+ const bool alias = overlaps(dst, x) || overlaps(dst, rope->src[1]) || (count > 1 && overlaps(dst, nodes[1]->src[1]));
+ void* output = alias ? ggml_backend_buffer_get_base(execution->workspace) : dst->data;
+ CUDAError error;
+ if (count > 1) {
+ error = rms_rope_launch(stream, nodes, output, count == 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<true> : rope_apply_f32<false>;
+ kernel<<<static_cast<unsigned int>((pairs + 255) / 256), 256, 0, stream>>>(static_cast<const float*>(x->data), static_cast<const float*>(pe->data),
+ static_cast<float*>(output), x->ne[0] / 2, pe->ne[2], x->ne[1], x->ne[2], pairs, strides);
+ error = cuda_last_error();
+ }
+ if (error == cuda_success && alias)
+ error = cuda_copy_async(dst->data, output, ggml_nbytes(dst), stream);
+ return error == cuda_success ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+
+ void register_rope(ggml_backend_dev_t device) {
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.rope_apply." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_rope_op();
+ desc.priority = 100;
+ desc.supports = rope_supports;
+ desc.prepare = rope_prepare;
+ desc.destroy = rope_destroy;
+ desc.execute = rope_execute;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ desc.name = "sd.rms_norm_mul_rope." SD_CUDA_NAME ".v1";
+ desc.op = GGML_OP_RMS_NORM;
+ desc.custom_op = nullptr;
+ desc.priority = 250;
+ desc.supports = nullptr;
+ desc.match = rope_match;
+ desc.optimize = rope_optimize;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ }
+
+}
diff --git a/src/core/cuda/rope.h b/src/core/cuda/rope.h
new file mode 100644
index 0000000..0d45f68
--- /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 0000000..f0db1dc
--- /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 <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+#define SD_CUDA_NAMESPACE sd_hip
+#define SD_CUDA_NAME "hip"
+#elif defined(SD_CUDA_VENDOR_MUSA)
+#include <musa_fp16.h>
+#include <musa_runtime.h>
+#define SD_CUDA_NAMESPACE sd_musa
+#define SD_CUDA_NAME "musa"
+#else
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#define SD_CUDA_NAMESPACE sd_cuda
+#define SD_CUDA_NAME "cuda"
+#endif
+
+namespace SD_CUDA_NAMESPACE {
+#if defined(SD_CUDA_VENDOR_HIP)
+ using CUDAError = hipError_t;
+ using CUDAStream = hipStream_t;
+ static constexpr auto cuda_success = hipSuccess;
+ static constexpr auto cuda_invalid_value = hipErrorInvalidValue;
+ static constexpr auto cuda_kind = GGML_EXT_GPU_HIP;
+ static CUDAError cuda_last_error() {
+ return hipGetLastError();
+ }
+ static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+ return hipMemsetAsync(dst, value, bytes, stream);
+ }
+ static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+ return hipMemcpyAsync(dst, src, bytes, hipMemcpyDeviceToDevice, stream);
+ }
+ template <typename Kernel>
+ static CUDAError cuda_set_shared_memory(Kernel, int) {
+ // HIP uses the device's shared-memory limit without CUDA's opt-in attribute.
+ return cuda_success;
+ }
+#elif defined(SD_CUDA_VENDOR_MUSA)
+ using CUDAError = musaError_t;
+ using CUDAStream = musaStream_t;
+ static constexpr auto cuda_success = musaSuccess;
+ static constexpr auto cuda_invalid_value = musaErrorInvalidValue;
+ static constexpr auto cuda_kind = GGML_EXT_GPU_MUSA;
+ static CUDAError cuda_last_error() {
+ return musaGetLastError();
+ }
+ static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+ return musaMemsetAsync(dst, value, bytes, stream);
+ }
+ static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+ return musaMemcpyAsync(dst, src, bytes, musaMemcpyDeviceToDevice, stream);
+ }
+ template <typename Kernel>
+ static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+ return musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+ }
+#else
+ using CUDAError = cudaError_t;
+ using CUDAStream = cudaStream_t;
+ static constexpr auto cuda_success = cudaSuccess;
+ static constexpr auto cuda_invalid_value = cudaErrorInvalidValue;
+ static constexpr auto cuda_kind = GGML_EXT_GPU_CUDA;
+ static CUDAError cuda_last_error() {
+ return cudaGetLastError();
+ }
+ static CUDAError cuda_memset_async(void* dst, int value, size_t bytes, CUDAStream stream) {
+ return cudaMemsetAsync(dst, value, bytes, stream);
+ }
+ static CUDAError cuda_copy_async(void* dst, const void* src, size_t bytes, CUDAStream stream) {
+ return cudaMemcpyAsync(dst, src, bytes, cudaMemcpyDeviceToDevice, stream);
+ }
+ template <typename Kernel>
+ static CUDAError cuda_set_shared_memory(Kernel kernel, int bytes) {
+ return cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, bytes);
+ }
+#endif
+ template <typename T>
+ static __device__ __forceinline__ T cuda_shuffle(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+ return __shfl(value, lane, width);
+#else
+ return __shfl_sync(0xffffffff, value, lane, width);
+#endif
+ }
+ template <typename T>
+ static __device__ __forceinline__ T cuda_shuffle_xor(T value, int lane, int width = 32) {
+#if defined(SD_CUDA_VENDOR_HIP)
+ return __shfl_xor(value, lane, width);
+#else
+ return __shfl_xor_sync(0xffffffff, value, lane, width);
+#endif
+ }
+ static bool cuda_context(ggml_backend_t backend, ggml_ext_gpu_context* context) {
+ auto reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend));
+ auto query = reinterpret_cast<ggml_ext_gpu_context_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_gpu_context_v1"));
+ return query && query(backend, context) && context->kind == cuda_kind;
+ }
+ static ggml_ext_gpu_device_info cuda_device_info(ggml_backend_dev_t device) {
+ ggml_ext_gpu_device_info info = {};
+ auto query = reinterpret_cast<ggml_ext_gpu_device_fn>(ggml_backend_reg_get_proc_address(ggml_backend_dev_backend_reg(device), "ggml_ext_gpu_device_v1"));
+ if (!query || !query(device, &info) || info.kind != cuda_kind)
+ return {cuda_kind, -1, 0, 0, false};
+ return info;
+ }
+}
+#endif
diff --git a/src/core/cuda/sage_attention.cu b/src/core/cuda/sage_attention.cu
new file mode 100644
index 0000000..bb42d5b
--- /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 <algorithm>
+#include <climits>
+#include <cstring>
+#include <map>
+#include "sageattention/attn.cuh"
+#include "sageattention/attn_fp8.cuh"
+
+static constexpr int sage_cuda_version = __CUDACC_VER_MAJOR__ * 1000 + __CUDACC_VER_MINOR__ * 10;
+
+namespace sd_cuda {
+
+ struct SageDevice {
+ int cc;
+ int compiled_cc;
+ };
+
+ static ggml_sage_attn_mode sage_mode(const SageDevice& device, ggml_sage_attn_mode requested) {
+ if (requested != GGML_SAGE_ATTN_AUTO)
+ return requested;
+ if (sage_cuda_version >= 12040 && device.cc >= 89 && device.compiled_cc >= 89) {
+ if (sage_cuda_version >= 12080 && device.cc != 90)
+ return GGML_SAGE_ATTN_2_PLUS_PLUS;
+ return GGML_SAGE_ATTN_2;
+ }
+ return GGML_SAGE_ATTN_AUTO;
+ }
+
+ static __device__ __forceinline__ float warp_reduce_max(float x) {
+#pragma unroll
+ for (int offset = 16; offset > 0; offset >>= 1)
+ x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, offset));
+ return x;
+ }
+
+ static __global__ void sage_output(const half* input, float* output, int count) {
+ const int i = blockIdx.x * blockDim.x + threadIdx.x;
+ if (i < count)
+ output[i] = __half2float(input[i]);
+ }
+
+ struct SageWorkspace {
+ size_t q;
+ size_t k;
+ size_t q_scale;
+ size_t k_scale;
+ size_t k_mean;
+ size_t v;
+ size_t v_scale;
+ size_t output;
+ size_t size;
+
+ explicit SageWorkspace(const ggml_tensor* dst) {
+ const ggml_tensor* qt = dst->src[0];
+ const ggml_tensor* kt = dst->src[1];
+ size = 0;
+ q = reserve(ggml_nelements(qt));
+ k = reserve(ggml_nelements(kt));
+ q_scale = reserve(GGML_PAD(qt->ne[1], 128) / 32 * 8 * qt->ne[2] * qt->ne[3] * sizeof(float));
+ k_scale = reserve(GGML_PAD(kt->ne[1], 64) / 64 * 4 * kt->ne[2] * kt->ne[3] * sizeof(float));
+ k_mean = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
+ v = reserve(kt->ne[0] * GGML_PAD(kt->ne[1], 64) * kt->ne[2] * kt->ne[3]);
+ v_scale = reserve(kt->ne[0] * kt->ne[2] * kt->ne[3] * sizeof(float));
+ output = reserve(ggml_nelements(dst) * sizeof(half));
+ }
+
+ size_t reserve(size_t bytes) {
+ const size_t offset = size;
+ size += GGML_PAD(bytes, 256);
+ return offset;
+ }
+ };
+
+ template <int D>
+ static __global__ void sage_key_mean(const float* k, float* mean, int length) {
+ const int channel = blockIdx.x * 32 + threadIdx.x % 32;
+ const int row = threadIdx.x / 32;
+ const size_t base = size_t(blockIdx.y) * length * D;
+ float sum = 0.0f;
+ for (int token = row; token < length; token += 8) {
+ sum += k[base + size_t(token) * D + channel];
+ }
+ __shared__ float sums[256];
+ sums[threadIdx.x] = sum;
+ __syncthreads();
+ if (row == 0) {
+ for (int i = 1; i < 8; ++i) {
+ sum += sums[threadIdx.x + i * 32];
+ }
+ mean[size_t(blockIdx.y) * D + channel] = sum / length;
+ }
+ }
+
+ template <int D, int TOKENS, bool SMOOTH>
+ static __global__ void sage_quantize(const float* input, const float* mean, int8_t* output, float* scales, int length) {
+ constexpr int VALUES = D * TOKENS / 256;
+ const int head = blockIdx.y;
+ const int first = blockIdx.x * TOKENS * D;
+ const size_t base = size_t(head) * length * D;
+ float values[VALUES];
+ float amax = 1e-7f;
+#pragma unroll
+ for (int i = 0; i < VALUES; ++i) {
+ const int index = first + threadIdx.x + i * 256;
+ float value = 0.0f;
+ if (index < length * D) {
+ value = input[base + index];
+ if constexpr (SMOOTH) {
+ value -= mean[size_t(head) * D + index % D];
+ }
+ }
+ values[i] = value;
+ amax = fmaxf(amax, fabsf(value));
+ }
+ amax = warp_reduce_max(amax);
+ __shared__ float maxima[8];
+ if (threadIdx.x % 32 == 0) {
+ maxima[threadIdx.x / 32] = amax;
+ }
+ __syncthreads();
+ if (threadIdx.x < 32) {
+ amax = threadIdx.x < 8 ? maxima[threadIdx.x] : 1e-7f;
+ amax = warp_reduce_max(amax);
+ if (threadIdx.x == 0) {
+ maxima[0] = amax;
+ scales[size_t(head) * gridDim.x + blockIdx.x] = amax / 127.0f;
+ }
+ }
+ __syncthreads();
+ const float inverse_scale = 127.0f / maxima[0];
+#pragma unroll
+ for (int i = 0; i < VALUES; ++i) {
+ const int index = first + threadIdx.x + i * 256;
+ if (index < length * D) {
+ output[base + index] = int8_t(__float2int_rn(values[i] * inverse_scale));
+ }
+ }
+ }
+
+ template <int D, bool KEY>
+ static __global__ void sage_quantize_per_thread(const float* input, const float* mean, int8_t* output, float* scales, int length) {
+ constexpr int GROUPS = KEY ? 4 : 8;
+ constexpr int TOKENS = KEY ? 64 : 32;
+ constexpr int VALUES = D * TOKENS / GROUPS / 32;
+ const int group = threadIdx.x / 32;
+ const int lane = threadIdx.x % 32;
+ const size_t base = size_t(blockIdx.y) * length * D;
+ float values[VALUES];
+ float amax = 0.0f;
+#pragma unroll
+ for (int i = 0; i < VALUES; ++i) {
+ const int index = lane + i * 32;
+ const int row = index / D;
+ const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
+ float value = token < length ? input[base + size_t(token) * D + index % D] : 0.0f;
+ if constexpr (KEY) {
+ if (token < length) {
+ value -= mean[size_t(blockIdx.y) * D + index % D];
+ }
+ }
+ values[i] = value;
+ amax = fmaxf(amax, fabsf(value));
+ }
+ const float scale = warp_reduce_max(amax) / 127.0f + 1e-7f;
+ if (lane == 0) {
+ scales[(size_t(blockIdx.y) * gridDim.x + blockIdx.x) * GROUPS + group] = scale;
+ }
+#pragma unroll
+ for (int i = 0; i < VALUES; ++i) {
+ const int index = lane + i * 32;
+ const int row = index / D;
+ const int token = blockIdx.x * TOKENS + (KEY ? row / 2 * 8 + group * 2 + row % 2 : row * 8 + group);
+ if (token < length) {
+ const float value = values[i] / scale;
+ output[base + size_t(token) * D + index % D] = int8_t(__float2int_rz(value + copysignf(0.5f, value)));
+ }
+ }
+ }
+
+ template <int D>
+ static __global__ void sage_value_scale(const half* v, float* scales, int length, float scale_max) {
+ const int channel = blockIdx.x * 32 + threadIdx.x % 32;
+ const int row = threadIdx.x / 32;
+ const size_t base = size_t(blockIdx.y) * length * D;
+ float amax = 1e-7f;
+ for (int token = row; token < length; token += 8) {
+ amax = fmaxf(amax, fabsf(__half2float(v[base + size_t(token) * D + channel])));
+ }
+ __shared__ float maxima[256];
+ maxima[threadIdx.x] = amax;
+ __syncthreads();
+ if (row == 0) {
+ for (int i = 1; i < 8; ++i) {
+ amax = fmaxf(amax, maxima[threadIdx.x + i * 32]);
+ }
+ scales[size_t(blockIdx.y) * D + channel] = amax / scale_max;
+ }
+ }
+
+ template <int D>
+ static __global__ void sage_quantize_value(const half* v, const float* scales, int8_t* output, int length, int padded) {
+ __shared__ uint8_t tile[32][D + 4];
+ const int head = blockIdx.y;
+ const int first = blockIdx.x * 32;
+ const size_t base = size_t(head) * length * D;
+ for (int i = threadIdx.x; i < 32 * D; i += 256) {
+ const int token = first + i / D;
+ const int channel = i % D;
+ const float value = token < length ? __half2float(v[base + size_t(token) * D + channel]) / scales[size_t(head) * D + channel] : 0.0f;
+ // FP8 MMA consumes pairs of tokens in a permuted, transposed V tile.
+ const int row = i / D;
+ const int permuted = row / 16 * 16 + row % 16 / 8 * 2 + row % 8 / 2 * 4 + row % 2;
+ tile[permuted][channel] = __nv_cvt_float_to_fp8(value, __NV_SATFINITE, __NV_E4M3);
+ }
+ __syncthreads();
+ for (int i = threadIdx.x; i < 32 * D; i += 256) {
+ output[(size_t(head) * D + i / 32) * padded + first + i % 32] = tile[i % 32][i / 32];
+ }
+ }
+
+ template <int D, bool PLUS_PLUS>
+ static void sage_attn_fp8(cudaStream_t stream, ggml_tensor* dst, char* data, const SageWorkspace& work, float scale) {
+ const ggml_tensor* q = dst->src[0];
+ const ggml_tensor* k = dst->src[1];
+ const ggml_tensor* v = dst->src[2];
+ const int lq = q->ne[1], lk = k->ne[1], padded = GGML_PAD(lk, 64);
+ const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
+ auto* qi = reinterpret_cast<int8_t*>(data + work.q);
+ auto* ki = reinterpret_cast<int8_t*>(data + work.k);
+ auto* vi = reinterpret_cast<int8_t*>(data + work.v);
+ auto* qs = reinterpret_cast<float*>(data + work.q_scale);
+ auto* ks = reinterpret_cast<float*>(data + work.k_scale);
+ auto* vs = reinterpret_cast<float*>(data + work.v_scale);
+ auto* km = reinterpret_cast<float*>(data + work.k_mean);
+ auto* out = reinterpret_cast<half*>(data + work.output);
+ sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
+ sage_quantize_per_thread<D, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
+ static_cast<const float*>(q->data), nullptr, qi, qs, lq);
+ sage_quantize_per_thread<D, true><<<dim3(padded / 64, hk * batch), 128, 0, stream>>>(
+ static_cast<const float*>(k->data), km, ki, ks, lk);
+ // The smaller V range keeps the FP16 instruction accumulator finite in SageAttention2++.
+ sage_value_scale<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, lk, PLUS_PLUS ? 2.25f : 448.0f);
+ sage_quantize_value<D><<<dim3(padded / 32, hk * batch), 256, 0, stream>>>(static_cast<const half*>(v->data), vs, vi, lk, padded);
+ auto kernel = sageattention::qk_int_sv_f8_attn_kernel<128, 64, 32, 64, D,
+ sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerThread,
+ sageattention::QuantGranularity::kPerThread, float, true, half,
+ sageattention::ComputeUnit::kCudaCore, sageattention::MaskMode::kNone, false, true, false, PLUS_PLUS>;
+ kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), 256 * D, stream>>>(
+ qi, ki, vi, out, nullptr, qs, ks, vs, nullptr, lq, lk, hq / hk,
+ D * lq * hq, D, D * lq, D * lk * hk, D, D * lk,
+ D * padded * hk, D * padded, padded, D * lq * hq, D * hq, D, scale);
+ sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
+ }
+
+ template <int D>
+ static void sage_attn(cudaStream_t stream, ggml_tensor* dst, char* data, ggml_sage_attn_mode mode, float scale) {
+ const ggml_tensor* q = dst->src[0];
+ const ggml_tensor* k = dst->src[1];
+ const ggml_tensor* v = dst->src[2];
+ const int lq = q->ne[1], lk = k->ne[1];
+ const int hq = q->ne[2], hk = k->ne[2], batch = q->ne[3];
+ const SageWorkspace work(dst);
+ if (mode == GGML_SAGE_ATTN_2_PLUS_PLUS) {
+ sage_attn_fp8<D, true>(stream, dst, data, work, scale);
+ return;
+ }
+ if (mode == GGML_SAGE_ATTN_2) {
+ sage_attn_fp8<D, false>(stream, dst, data, work, scale);
+ return;
+ }
+ auto* qi = reinterpret_cast<int8_t*>(data + work.q);
+ auto* ki = reinterpret_cast<int8_t*>(data + work.k);
+ auto* qs = reinterpret_cast<float*>(data + work.q_scale);
+ auto* ks = reinterpret_cast<float*>(data + work.k_scale);
+ auto* km = reinterpret_cast<float*>(data + work.k_mean);
+ auto* out = reinterpret_cast<half*>(data + work.output);
+
+ sage_key_mean<D><<<dim3(D / 32, hk * batch), 256, 0, stream>>>(static_cast<const float*>(k->data), km, lk);
+ sage_quantize<D, 32, false><<<dim3(GGML_PAD(lq, 128) / 32, hq * batch), 256, 0, stream>>>(
+ static_cast<const float*>(q->data), nullptr, qi, qs, lq);
+ sage_quantize<D, 64, true><<<dim3(GGML_PAD(lk, 64) / 64, hk * batch), 256, 0, stream>>>(
+ static_cast<const float*>(k->data), km, ki, ks, lk);
+
+ constexpr size_t shared_bytes = (128 + 64 + 64 * 2) * D;
+ auto kernel = sageattention::qk_int_sv_f16_attn_kernel<128, 64, 32, 64, D,
+ sageattention::DataType::kInt8, sageattention::QuantGranularity::kPerWarp,
+ sageattention::QuantGranularity::kPerWarp, float, false, half,
+ sageattention::ComputeUnit::kTensorCore, sageattention::MaskMode::kNone, false, false>;
+ kernel<<<dim3(GGML_PAD(lq, 128) / 128, hq, batch), dim3(32, 4), shared_bytes, stream>>>(
+ qi, ki, static_cast<half*>(v->data), out, nullptr, qs, ks, nullptr,
+ lq, lk, hq / hk,
+ D * lq * hq, D, D * lq,
+ D * lk * hk, D, D * lk,
+ D * lk * hk, D, D * lk,
+ D * lq * hq, D * hq, D,
+ scale);
+ sage_output<<<(ggml_nelements(dst) + 255) / 256, 256, 0, stream>>>(out, static_cast<float*>(dst->data), ggml_nelements(dst));
+ }
+
+ static bool sage_supports(ggml_backend_dev_t, const ggml_tensor* op, void* user_data) {
+ const auto& device = *static_cast<const SageDevice*>(user_data);
+ SDSageAttentionParams params;
+ if (!sd_ggml_sage_attn_describe(op, ¶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<uintptr_t>(t->data) & 15))
+ return false;
+ }
+ return k->ne[0] * GGML_PAD(k->ne[1], 64) * k->ne[2] * k->ne[3] <= INT_MAX;
+ }
+
+ static ggml_status sage_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+ *workspace = SageWorkspace(nodes[0]).size;
+ return GGML_STATUS_SUCCESS;
+ }
+
+ static void sage_destroy(void*, void*) {}
+
+ static ggml_status sage_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void* user_data) {
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ SDSageAttentionParams params;
+ if (!sd_ggml_sage_attn_describe(nodes[0], ¶ms))
+ return GGML_STATUS_FAILED;
+ const auto mode = sage_mode(*static_cast<const SageDevice*>(user_data), params.mode);
+ auto* data = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ auto stream = static_cast<cudaStream_t>(context.stream);
+ if (nodes[0]->src[0]->ne[0] == 64)
+ sage_attn<64>(stream, nodes[0], data, mode, params.scale);
+ else
+ sage_attn<128>(stream, nodes[0], data, mode, params.scale);
+ return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+
+ void register_sage_attention(ggml_backend_dev_t device) {
+ const auto info = cuda_device_info(device);
+ if (info.device < 0)
+ return;
+ int previous;
+ cudaDeviceProp properties;
+ cudaFuncAttributes attributes;
+ if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || cudaGetDevice(&previous) != cudaSuccess ||
+ cudaSetDevice(info.device) != cudaSuccess)
+ return;
+ const auto error = cudaFuncGetAttributes(&attributes, sage_output);
+ if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess)
+ return;
+ // PTX compiled for an older target retains that target's instruction guards after JIT.
+ static std::map<int, SageDevice> devices;
+ auto& data = devices[info.device];
+ data = {properties.major * 10 + properties.minor, std::min(attributes.binaryVersion, attributes.ptxVersion)};
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.sage_attn.cuda.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_sage_attn_op();
+ desc.priority = 100;
+ desc.supports = sage_supports;
+ desc.prepare = sage_prepare;
+ desc.destroy = sage_destroy;
+ desc.execute = sage_execute;
+ desc.user_data = &data;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ }
+
+}
+#else
+namespace sd_cuda {
+ void register_sage_attention(ggml_backend_dev_t) {}
+}
+#endif
diff --git a/src/core/cuda/sageattention/LICENSE b/src/core/cuda/sageattention/LICENSE
new file mode 100644
index 0000000..927221d
--- /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 0000000..ea43a31
--- /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 `<cuda/pipeline>` includes were removed so multi-architecture builds can
+also compile targets below SM70; the backend rejects SageAttention below SM80.
+Predicated loads now zero unused FP16 V lanes: masked probabilities cannot
+suppress NaNs left in shared memory by previous kernels.
+
+The sd.cpp adapter in `../sage_attention.cu` supplies FP32 K smoothing and native CUDA
+per-thread INT8 quantization, matching the token grouping in upstream
+`sageattention/triton/quant_per_thread.py`. V is quantized per channel and
+transposed, padded with zeros, and permuted according to
+`csrc/fused/fused.cu`. Its FP8 range is 448 for SageAttention2 and 2.25 for
+SageAttention2++. Zero channels use a positive scale to avoid division by zero.
+SageAttention2 uses an FP32 instruction buffer; 2++ uses an FP16 instruction
+buffer added to FP32 accumulators. The SM80 compatibility path retains per-warp
+INT8 Q/K and FP16 V with FP32 accumulation.
+
+The adapter requests workspace and the current stream through public ggml extension APIs; it has no Python, PyTorch,
+or Triton dependency. The public output is converted back to FP32.
diff --git a/src/core/cuda/sageattention/attn.cuh b/src/core/cuda/sageattention/attn.cuh
new file mode 100644
index 0000000..540d582
--- /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 <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_mean = false>
+ __global__ void qk_int_sv_f16_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, half* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, DTypeOut* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_seq_v, const uint32_t stride_h_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
+ // Older CUDA headers omit BF16 intrinsics in pre-Ampere device passes.
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 800
+ // compile time check
+ static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4");
+ static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread");
+ static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread");
+ static_assert(std::is_same<DTypeSVAccum, float>::value || !use_inst_buffer, "use_inst_buffer only supports DTypeSVAccum as float");
+ static_assert(std::is_same<DTypeSVAccum, float>::value || std::is_same<DTypeSVAccum, half>::value, "DTypeSVAccum must be float or half");
+ static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::value, "DTypeOut must be half or nv_bfloat16");
+ static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64");
+ static_assert(!fuse_v_mean || std::is_same<DTypeSVAccum, half>::value, "fuse_v_mean only supports half");
+ static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation
+
+ using DTypeOut2 = typename std::conditional<std::is_same<DTypeOut, half>::value, half2, nv_bfloat162>::type;
+
+ constexpr uint32_t num_warps_q = CTA_Q / WARP_Q;
+ constexpr uint32_t num_warps_k = CTA_K / WARP_K;
+ constexpr uint32_t num_warps = num_warps_q * num_warps_k;
+ constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M;
+ constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N;
+ constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K);
+ constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N;
+
+ constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2);
+ constexpr uint32_t O_SMEM_STRIDE = head_dim;
+ constexpr uint32_t V_SMEM_STRIDE = head_dim;
+
+ extern __shared__ int8_t smem[];
+
+ const uint32_t lane_id = get_lane_id();
+ const uint32_t warp_id = get_warp_id();
+
+ // maximize L2 hit rate
+ const uint32_t batch_id = blockIdx.z;
+ const uint32_t bx = blockIdx.x;
+ const uint32_t num_qo_heads = gridDim.y;
+ const uint32_t head_id = blockIdx.y;
+
+ // transfer to base 2 instead of base e with better numerical efficiency
+ sm_scale *= math::log2e;
+
+ // RS holds the fragment of S
+ int32_t RS[num_tiles_q][num_tiles_k][8];
+ DTypeSVAccum RO[num_tiles_q][num_tiles_v][8];
+ float m[num_tiles_q][2]; // max
+ float d[num_tiles_q][2]; // denominator
+
+ uint32_t q_scale_idx, k_scale_idx;
+
+ if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
+ const uint32_t num_block_q = gridDim.x;
+ q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
+ } else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
+ const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+ q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
+ } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
+ const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+ q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
+ }
+
+ if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
+ const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
+ k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
+ } else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
+ const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+ k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
+ } else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
+ const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+ k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
+ }
+
+ constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
+ : (CTA_K / WARP_K) * 4;
+
+ // initialize o, m, d
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RO[fq][fv][k] = 0.0f;
+ }
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+#pragma unroll
+ for (uint32_t k = 0; k < 4; k++) {
+ ((int32_t*)RO[fq][fv])[k] = 0;
+ }
+ }
+ }
+ }
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 2; k++) {
+ m[fq][k] = -5000000.0f;
+ d[fq][k] = 1.0f;
+ }
+ }
+
+ constexpr uint32_t K_smem_idx_offset = CTA_Q;
+ constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
+
+ constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
+ : SwizzleMode::k128B;
+ smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
+ smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
+ constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
+ smem_t<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE);
+ constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
+ smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
+
+ constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
+ : 8;
+ constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
+ : 4;
+ constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 32) ? 4 : 8;
+ constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 32) ? 8 : 4;
+ constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8;
+ constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4;
+
+ constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK);
+ constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+ constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+ constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V);
+ constexpr uint32_t V_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_V);
+ constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
+ constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
+
+ int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+ int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+ half* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V) * stride_seq_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
+ uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+ uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+ uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
+
+ uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
+ uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2);
+ uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 16, lane_id / 16);
+
+ // for causal masking
+ uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+ uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + 2 * (lane_id % 4);
+
+ // for loading
+ uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+ uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+ uint32_t V_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_V;
+
+ const uint32_t num_iterations = SAGE_DIV_CEIL(
+ mask_mode == MaskMode::kCausal
+ ? min(kv_len, (bx + 1) * CTA_Q)
+ : kv_len,
+ CTA_K);
+
+ // load Q with predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
+ &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len);
+ cp_async::commit_group();
+ cp_async::wait_group<0>();
+ __syncthreads();
+
+ // for num_tiles_qk_inner = 1, we load all Qs in register
+ uint32_t RQ[num_tiles_q][4];
+ if constexpr (num_tiles_qk_inner == 1) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
+ Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
+ }
+ }
+
+ // load K with predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+ cp_async::commit_group();
+
+ float q_scale = Q_scale[q_scale_idx];
+
+ float original_sm_scale = sm_scale;
+ float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset];
+
+ sm_scale = original_sm_scale * dequant_scale;
+
+ // load V with predicate
+ load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len);
+ cp_async::commit_group();
+
+ K_load_idx_lane_base += CTA_K;
+ V_load_idx_lane_base += CTA_K;
+
+#pragma unroll
+ for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
+ // ensure K is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // compute QK^T
+ if constexpr (num_tiles_qk_inner == 1) {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_K, RS, RQ, K_smem_offset_mma);
+ } else {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+ }
+
+ float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
+ }
+ }
+ }
+
+ // do not apply causal mask and out of bound mask for these iterations
+ K_idx_lane_base += CTA_K;
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, sm_scale);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, sm_scale);
+ }
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+
+ uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+ RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+ }
+
+ __syncthreads();
+
+ // load K
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K);
+ cp_async::commit_group();
+
+ dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset];
+ sm_scale = original_sm_scale * dequant_scale;
+
+ // ensure V is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ if constexpr (!use_inst_buffer) {
+ compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ } else {
+ compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ }
+
+ __syncthreads();
+ // load V
+ load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V);
+ cp_async::commit_group();
+ K_load_idx_lane_base += CTA_K;
+ V_load_idx_lane_base += CTA_K;
+ }
+
+ // second last iter, apply causal mask
+ if (num_iterations > 1) {
+ // ensure K is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // compute QK^T
+ if constexpr (num_tiles_qk_inner == 1) {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_K, RS, RQ, K_smem_offset_mma);
+ } else {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+ }
+
+ float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+ }
+ }
+ }
+
+ if constexpr (mask_mode == MaskMode::kCausal) {
+ apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+ }
+ // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+ K_idx_lane_base += CTA_K;
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
+ }
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+
+ uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+ RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+ }
+
+ __syncthreads();
+
+ // load K with predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+ cp_async::commit_group();
+
+ dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset];
+ sm_scale = original_sm_scale * dequant_scale;
+
+ // ensure V is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ if constexpr (!use_inst_buffer) {
+ compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ } else {
+ compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ }
+
+ __syncthreads();
+ // load V with predicate
+ load_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &V_lane_base_ptr, V_smem_offset_load, stride_seq_v, smem_V, V_load_idx_lane_base, kv_len);
+ cp_async::commit_group();
+ K_load_idx_lane_base += CTA_K;
+ V_load_idx_lane_base += CTA_K;
+ }
+
+ // last iter, apply causal mask and out of bound mask
+ {
+ // ensure K is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // compute QK^T
+ if constexpr (num_tiles_qk_inner == 1) {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_K, RS, RQ, K_smem_offset_mma);
+ } else {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+ }
+
+ float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+ }
+ }
+ }
+
+ if constexpr (mask_mode == MaskMode::kCausal) {
+ apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+ }
+ // check out of bound in the last iter
+ apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+ K_idx_lane_base += CTA_K;
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, false, false>(RS_f32, RO, m, d, original_sm_scale);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, false, false>(RS_f32, RO, m, d, original_sm_scale);
+ }
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+
+ uint32_t RS_f16[num_tiles_q][num_tiles_k][4];
+ RS_32_to_16<num_tiles_q, num_tiles_k>(RS_f32, RS_f16);
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kTensorCore>(RS_f16, d);
+ }
+
+ // ensure V is ready
+ cp_async::wait_group<0>();
+ __syncthreads();
+
+ if constexpr (!use_inst_buffer) {
+ compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ } else {
+ compute_fp16_sv_permuted_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ }
+
+ __syncthreads();
+ }
+
+ // TODO: thread block sync mdo state for num_warps_k > 0
+
+ normalize_d<num_tiles_q, num_tiles_v, DenominatorAccumUnit>(RO, m, d);
+
+ // save the result
+ // if (get_warp_idx_k<num_warps_q, num_warps_k>() == 0)
+ // {
+
+ // convert half to bfloat16
+ if constexpr (std::is_same<DTypeSVAccum, half>::value && std::is_same<DTypeOut, nv_bfloat16>::value) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ ((nv_bfloat162*)RO[fq][fv])[0] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[0]));
+ ((nv_bfloat162*)RO[fq][fv])[1] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[1]));
+ ((nv_bfloat162*)RO[fq][fv])[2] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[2]));
+ ((nv_bfloat162*)RO[fq][fv])[3] = __float22bfloat162_rn(__half22float2(((half2*)RO[fq][fv])[3]));
+ }
+ }
+ }
+
+ // add v_mean
+ if constexpr (fuse_v_mean) {
+ DTypeOut2 v_mean[2];
+ DTypeOut* V_mean_lane_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + lane_id % 4 * 2;
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ v_mean[0] = *((DTypeOut2*)(V_mean_lane_ptr + fv * 16));
+ v_mean[1] = *((DTypeOut2*)(V_mean_lane_ptr + 8 + fv * 16));
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ ((DTypeOut2*)RO[fq][fv])[0] = __hadd2(((DTypeOut2*)RO[fq][fv])[0], v_mean[0]);
+ ((DTypeOut2*)RO[fq][fv])[1] = __hadd2(((DTypeOut2*)RO[fq][fv])[1], v_mean[0]);
+ ((DTypeOut2*)RO[fq][fv])[2] = __hadd2(((DTypeOut2*)RO[fq][fv])[2], v_mean[1]);
+ ((DTypeOut2*)RO[fq][fv])[3] = __hadd2(((DTypeOut2*)RO[fq][fv])[3], v_mean[1]);
+ }
+ }
+ }
+
+ // save the result to shared memory
+ uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ // convert RO to half
+ uint32_t RO_f16[4];
+#pragma unroll
+ for (uint32_t k = 0; k < 4; k++) {
+ if constexpr (std::is_same<DTypeOut, half>::value) {
+ ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
+ } else if constexpr (std::is_same<DTypeOut, nv_bfloat16>::value) {
+ ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
+ }
+ }
+
+ ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0];
+ ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1];
+
+ // ! permuted, make sure you know what you are doing
+ ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = RO_f16[2];
+ ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[0];
+ ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[1];
+
+ // ! permuted, make sure you know what you are doing
+ ((uint32_t*)(smem_O.base + (offset_O ^ 0x1)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[2];
+ ((uint32_t*)(smem_O.base + (offset_O ^ 0x1) + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = ((uint32_t*)RO[fq][fv])[3];
+ }
+ }
+ }
+
+ // ! do we need to sync here?
+ __syncwarp();
+
+ // shared memory to global memory
+ DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
+ uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O);
+ uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
+
+#pragma unroll
+ for (uint32_t i = 0; i < O_smem_iters_col; i++) {
+#pragma unroll
+ for (uint32_t j = 0; j < O_smem_iters_row; j++) {
+ if (O_load_idx_lane_base < qo_len) {
+ smem_O.store_128b(offset_O, O_lane_ptr);
+ }
+ O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
+ offset_O = smem_O.advance_offset_by_column<global_to_shared_line_lanes_O>(offset_O);
+ }
+
+ offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O));
+ O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O));
+ O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
+ }
+
+ if constexpr (return_lse) {
+ uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
+ float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
+ uint32_t fq = (lane_id % 4) / 2;
+ uint32_t k = (lane_id % 4) % 2;
+
+ if (lse_idx < qo_len && (lane_id % 4) < 2 * num_tiles_q) {
+ lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
+ }
+ }
+
+ // }
+#else
+ RUNTIME_ASSERT("SageAttention requires SM80 or newer");
+#endif
+ }
+
+#undef PACK_SIZE_QK
+#undef PACK_SIZE_V
+#undef PACK_SIZE_O
+#undef MMA_QK_M
+#undef MMA_QK_N
+#undef MMA_QK_K
+#undef MMA_SV_M
+#undef MMA_SV_N
+#undef MMA_SV_K
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/attn_fp8.cuh b/src/core/cuda/sageattention/attn_fp8.cuh
new file mode 100644
index 0000000..838e4c6
--- /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 <uint32_t CTA_Q, uint32_t CTA_K, uint32_t WARP_Q, uint32_t WARP_K, uint32_t head_dim, DataType DTypeQK, QuantGranularity Q_GRAN, QuantGranularity K_GRAN, typename DTypeSVAccum = float, bool use_inst_buffer = false, typename DTypeOut = half, ComputeUnit DenominatorAccumUnit, MaskMode mask_mode = MaskMode::kNone, bool return_lse = false, bool fuse_v_scale = false, bool fuse_v_mean = false, bool use_pv_fp16_accu = false>
+ __global__ void qk_int_sv_f8_attn_kernel(int8_t* __restrict__ Q, int8_t* __restrict__ K, int8_t* __restrict__ V, DTypeOut* __restrict__ O, float* __restrict__ Lse, float* __restrict__ Q_scale, float* __restrict__ K_scale, float* __restrict__ V_scale, float* __restrict__ V_mean, const uint32_t qo_len, const uint32_t kv_len, const uint32_t num_kv_groups, const uint32_t stride_bz_q, const uint32_t stride_seq_q, const uint32_t stride_h_q, const uint32_t stride_bz_k, const uint32_t stride_seq_k, const uint32_t stride_h_k, const uint32_t stride_bz_v, const uint32_t stride_h_v, const uint32_t stride_d_v, const uint32_t stride_bz_o, const uint32_t stride_seq_o, const uint32_t stride_h_o, float sm_scale) {
+ // Keep unsupported device passes from parsing architecture-specific intrinsics.
+#if !defined(__CUDA_ARCH__) || __CUDA_ARCH__ >= 890
+ // compile time check
+ static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4, "DTypeQK must be int8 or int4");
+ static_assert(Q_GRAN == QuantGranularity::kPerBlock || Q_GRAN == QuantGranularity::kPerWarp || Q_GRAN == QuantGranularity::kPerThread, "Q_GRAN must be kPerBlock, kPerWarp or kPerThread");
+ static_assert(K_GRAN == QuantGranularity::kPerBlock || K_GRAN == QuantGranularity::kPerWarp || K_GRAN == QuantGranularity::kPerThread, "K_GRAN must be kPerBlock, kPerWarp or kPerThread");
+ static_assert(head_dim % 64 == 0, "head_dim must be a multiple of 64");
+ static_assert(std::is_same<DTypeSVAccum, float>::value, "DTypeSVAccum must be float, half is WIP");
+ static_assert(std::is_same<DTypeOut, half>::value || std::is_same<DTypeOut, nv_bfloat16>::value, "DTypeOut must be half or nv_bfloat16");
+ static_assert(CTA_K % 64 == 0);
+ static_assert(CTA_Q / CTA_K <= 2); // for efficient causal implementation
+
+ constexpr uint32_t num_warps_q = CTA_Q / WARP_Q;
+ constexpr uint32_t num_warps_k = CTA_K / WARP_K;
+ constexpr uint32_t num_warps = num_warps_q * num_warps_k;
+ constexpr uint32_t num_tiles_q = WARP_Q / MMA_QK_M;
+ constexpr uint32_t num_tiles_k = WARP_K / MMA_QK_N;
+ constexpr uint32_t num_tiles_qk_inner = (DTypeQK == DataType::kInt8) ? (head_dim / MMA_QK_K) : (head_dim / 2 / MMA_QK_K);
+ constexpr uint32_t num_tiles_v = head_dim / MMA_SV_N;
+
+ constexpr uint32_t QK_SMEM_STRIDE = (DTypeQK == DataType::kInt8) ? (head_dim) : (head_dim / 2);
+ constexpr uint32_t O_SMEM_STRIDE = head_dim;
+ // for fp16: head_dim
+ constexpr uint32_t V_SMEM_STRIDE = CTA_K;
+
+ extern __shared__ int8_t smem[];
+
+ const uint32_t lane_id = get_lane_id();
+ const uint32_t warp_id = get_warp_id();
+
+ // maximize L2 hit rate
+ const uint32_t batch_id = blockIdx.z;
+ const uint32_t bx = blockIdx.x;
+ const uint32_t num_qo_heads = gridDim.y;
+ const uint32_t head_id = blockIdx.y;
+
+ // transfer to base 2 instead of base e with better numerical efficiency
+ sm_scale *= math::log2e;
+
+ // RS holds the fragment of S
+ int32_t RS[num_tiles_q][num_tiles_k][8];
+ DTypeSVAccum RO[num_tiles_q][num_tiles_v][8];
+ float m[num_tiles_q][2]; // max
+ float d[num_tiles_q][2]; // denominator
+
+ uint32_t q_scale_idx, k_scale_idx;
+
+ if constexpr (Q_GRAN == QuantGranularity::kPerBlock) {
+ const uint32_t num_block_q = gridDim.x;
+ q_scale_idx = batch_id * num_qo_heads * num_block_q + head_id * num_block_q + bx;
+ } else if constexpr (Q_GRAN == QuantGranularity::kPerWarp) {
+ const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+ q_scale_idx = batch_id * num_qo_heads * num_warp_block_q + head_id * num_warp_block_q + bx * num_warps_q + get_warp_idx_q<num_warps_q, num_warps_k>();
+ } else if constexpr (Q_GRAN == QuantGranularity::kPerThread) {
+ const uint32_t num_warp_block_q = gridDim.x * num_warps_q;
+ q_scale_idx = batch_id * num_qo_heads * (num_warp_block_q * 8) + head_id * (num_warp_block_q * 8) + bx * (num_warps_q * 8) + get_warp_idx_q<num_warps_q, num_warps_k>() * 8 + lane_id / 4;
+ }
+
+ if constexpr (K_GRAN == QuantGranularity::kPerBlock) {
+ const uint32_t num_block_k = SAGE_DIV_CEIL(kv_len, CTA_K);
+ k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_block_k + (head_id / num_kv_groups) * num_block_k;
+ } else if constexpr (K_GRAN == QuantGranularity::kPerWarp) {
+ const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+ k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * num_warp_block_k + (head_id / num_kv_groups) * num_warp_block_k + get_warp_idx_k<num_warps_q, num_warps_k>();
+ } else if constexpr (K_GRAN == QuantGranularity::kPerThread) {
+ const uint32_t num_warp_block_k = SAGE_DIV_CEIL(kv_len, CTA_K) * (CTA_K / WARP_K);
+ k_scale_idx = batch_id * (num_qo_heads / num_kv_groups) * (num_warp_block_k * 4) + (head_id / num_kv_groups) * (num_warp_block_k * 4) + get_warp_idx_k<num_warps_q, num_warps_k>() * 4 + lane_id % 4;
+ }
+
+ constexpr uint32_t k_scale_advance_offset = (K_GRAN == QuantGranularity::kPerBlock) ? 1 : (K_GRAN == QuantGranularity::kPerWarp) ? (CTA_K / WARP_K)
+ : (CTA_K / WARP_K) * 4;
+
+ // initialize o, m, d
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RO[fq][fv][k] = 0.0f;
+ }
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+#pragma unroll
+ for (uint32_t k = 0; k < 4; k++) {
+ ((int32_t*)RO[fq][fv])[k] = 0;
+ }
+ }
+ }
+ }
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 2; k++) {
+ m[fq][k] = -5000000.0f;
+ d[fq][k] = 1.0f;
+ }
+ }
+
+ constexpr uint32_t K_smem_idx_offset = CTA_Q;
+ constexpr uint32_t V_smem_idx_offset = CTA_Q + CTA_K;
+
+ constexpr SwizzleMode swizzle_mode_QK = (QK_SMEM_STRIDE == 32) ? SwizzleMode::k32B : (QK_SMEM_STRIDE == 64) ? SwizzleMode::k64B
+ : SwizzleMode::k128B;
+ smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_Q(smem);
+ smem_t<swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK> smem_K(smem + K_smem_idx_offset * QK_SMEM_STRIDE);
+ // for fp16: 32
+ constexpr SwizzleMode swizzle_mode_V = (V_SMEM_STRIDE == 64) ? SwizzleMode::k64B : SwizzleMode::k128B;
+ smem_t<swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V> smem_V(smem + V_smem_idx_offset * QK_SMEM_STRIDE);
+ constexpr SwizzleMode swizzle_mode_O = (O_SMEM_STRIDE == 32) ? SwizzleMode::k64B : SwizzleMode::k128B;
+ smem_t<swizzle_mode_O, O_SMEM_STRIDE / PACK_SIZE_O> smem_O(smem);
+
+ constexpr uint32_t global_to_shared_line_lanes_QK = (QK_SMEM_STRIDE == 32) ? 2 : (QK_SMEM_STRIDE == 64) ? 4
+ : 8;
+ constexpr uint32_t global_to_shared_copy_lines_per_warp_QK = (QK_SMEM_STRIDE == 32) ? 16 : (QK_SMEM_STRIDE == 64) ? 8
+ : 4;
+ // for fp16: 32
+ constexpr uint32_t global_to_shared_line_lanes_V = (V_SMEM_STRIDE == 64) ? 4 : 8;
+ // for fp16: 32
+ constexpr uint32_t global_to_shared_copy_lines_per_warp_V = (V_SMEM_STRIDE == 64) ? 8 : 4;
+ constexpr uint32_t global_to_shared_line_lanes_O = (O_SMEM_STRIDE == 32) ? 4 : 8;
+ constexpr uint32_t global_to_shared_copy_lines_per_warp_O = (O_SMEM_STRIDE == 32) ? 8 : 4;
+
+ constexpr uint32_t QK_smem_iters_row = QK_SMEM_STRIDE / (global_to_shared_line_lanes_QK * PACK_SIZE_QK);
+ constexpr uint32_t Q_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+ constexpr uint32_t K_smem_iters_col = CTA_K / (num_warps * global_to_shared_copy_lines_per_warp_QK);
+ constexpr uint32_t V_smem_iters_row = V_SMEM_STRIDE / (global_to_shared_line_lanes_V * PACK_SIZE_V);
+ // for fp16: CTA_K
+ constexpr uint32_t V_smem_iters_col = head_dim / (num_warps * global_to_shared_copy_lines_per_warp_V);
+ constexpr uint32_t O_smem_iters_row = O_SMEM_STRIDE / (global_to_shared_line_lanes_O * PACK_SIZE_O);
+ constexpr uint32_t O_smem_iters_col = CTA_Q / (num_warps * global_to_shared_copy_lines_per_warp_O);
+
+ int8_t* Q_lane_base_ptr = Q + batch_id * stride_bz_q + head_id * stride_h_q + (bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_q + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+ int8_t* K_lane_base_ptr = K + batch_id * stride_bz_k + (head_id / num_kv_groups) * stride_h_k + (CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK) * stride_seq_k + (lane_id % global_to_shared_line_lanes_QK) * PACK_SIZE_QK;
+ // for fp16: CTA_K / num_warps * warp_id * stride_seq_v + lane_id / global_to_shared_line_lanes_V * stride_seq_v
+ int8_t* V_lane_base_ptr = V + batch_id * stride_bz_v + (head_id / num_kv_groups) * stride_h_v + head_dim / num_warps * warp_id * stride_d_v + lane_id / global_to_shared_line_lanes_V * stride_d_v + (lane_id % global_to_shared_line_lanes_V) * PACK_SIZE_V;
+ uint32_t Q_smem_offset_load = smem_Q.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * Q_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+ uint32_t K_smem_offset_load = smem_K.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_QK * K_smem_iters_col + lane_id / global_to_shared_line_lanes_QK, lane_id % global_to_shared_line_lanes_QK);
+ uint32_t V_smem_offset_load = smem_V.get_permuted_offset(warp_id * global_to_shared_copy_lines_per_warp_V * V_smem_iters_col + lane_id / global_to_shared_line_lanes_V, lane_id % global_to_shared_line_lanes_V);
+
+ uint32_t Q_smem_offset_mma = smem_Q.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id % 16, lane_id / 16);
+ uint32_t K_smem_offset_mma = smem_K.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 8 + (lane_id / 16) * 8, (lane_id / 8) % 2);
+ // for fp 16:
+ // uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + lane_id % 16, lane_id / 16);
+ uint32_t V_smem_offset_mma = smem_V.get_permuted_offset(lane_id % 8 + (lane_id / 16) * 8, get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K / PACK_SIZE_V + (lane_id / 8) % 2);
+
+ // for causal masking
+ uint32_t Q_idx_lane_base = bx * CTA_Q + get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+ uint32_t K_idx_lane_base = get_warp_idx_k<num_warps_q, num_warps_k>() * WARP_K + 2 * (lane_id % 4);
+
+ // for loading
+ uint32_t Q_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+ uint32_t K_load_idx_lane_base = CTA_K / num_warps * warp_id + lane_id / global_to_shared_line_lanes_QK;
+
+ const uint32_t num_iterations = SAGE_DIV_CEIL(
+ mask_mode == MaskMode::kCausal
+ ? min(kv_len, (bx + 1) * CTA_Q)
+ : kv_len,
+ CTA_K);
+
+ // load Q with predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, Q_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_Q>(
+ &Q_lane_base_ptr, Q_smem_offset_load, stride_seq_q, smem_Q, Q_load_idx_lane_base, qo_len);
+ cp_async::commit_group();
+ cp_async::wait_group<0>();
+ __syncthreads();
+
+ // for num_tiles_qk_inner = 1, we load all Qs in register
+ uint32_t RQ[num_tiles_q][4];
+ if constexpr (num_tiles_qk_inner == 1) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ smem_Q.ldmatrix_m8n8x4(Q_smem_offset_mma, RQ[fq]);
+ Q_smem_offset_mma = smem_Q.advance_offset_by_row<16>(Q_smem_offset_mma);
+ }
+ }
+
+ // load K with predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+ cp_async::commit_group();
+
+ float q_scale = Q_scale[q_scale_idx];
+
+ float original_sm_scale = sm_scale;
+ float dequant_scale = q_scale * K_scale[k_scale_idx + 0 * k_scale_advance_offset];
+
+ sm_scale = original_sm_scale * dequant_scale;
+
+ // load V
+ // ! we assume that V is padded. If not, there might be illegal memory access or nan issue.
+ // for fp16:
+ // load_global_to_share stride_seq_v
+ load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
+ cp_async::commit_group();
+
+ K_load_idx_lane_base += CTA_K;
+
+#pragma unroll
+ for (uint32_t iter = 1; iter < num_iterations - 1; iter++) {
+ // ensure K is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // compute QK^T
+ if constexpr (num_tiles_qk_inner == 1) {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_K, RS, RQ, K_smem_offset_mma);
+ } else {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+ }
+ float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]);
+ }
+ }
+ }
+
+ K_idx_lane_base += CTA_K;
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, sm_scale);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, sm_scale);
+ }
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+
+ uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+ RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+ accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+ }
+
+ __syncthreads();
+
+ // load K without predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K);
+ cp_async::commit_group();
+
+ dequant_scale = q_scale * K_scale[k_scale_idx + iter * k_scale_advance_offset];
+ sm_scale = original_sm_scale * dequant_scale;
+
+ // ensure V is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // for fp16:
+ // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ // smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ if constexpr (!use_inst_buffer) {
+ compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ } else {
+ if constexpr (!use_pv_fp16_accu) {
+ compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ } else {
+ compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ }
+ }
+ __syncthreads();
+ // load V
+ // for fp16:
+ // load_global_to_share stride_seq_v
+ load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
+ cp_async::commit_group();
+
+ K_load_idx_lane_base += CTA_K;
+ }
+
+ // second last iter, apply causal mask
+ if (num_iterations > 1) {
+ // ensure K is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // compute QK^T
+ if constexpr (num_tiles_qk_inner == 1) {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_K, RS, RQ, K_smem_offset_mma);
+ } else {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+ }
+
+ float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+ }
+ }
+ }
+
+ if constexpr (mask_mode == MaskMode::kCausal) {
+ apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+ }
+ // apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+ K_idx_lane_base += CTA_K;
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
+ }
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+
+ uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+ RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+ accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+ }
+
+ __syncthreads();
+
+ // load K with predicate
+ load_global_to_share<global_to_shared_line_lanes_QK, global_to_shared_copy_lines_per_warp_QK, QK_smem_iters_row, K_smem_iters_col, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, CTA_K>(
+ &K_lane_base_ptr, K_smem_offset_load, stride_seq_k, smem_K, K_load_idx_lane_base, kv_len);
+ cp_async::commit_group();
+
+ dequant_scale = q_scale * K_scale[k_scale_idx + (num_iterations - 1) * k_scale_advance_offset];
+ sm_scale = original_sm_scale * dequant_scale;
+
+ // ensure V is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // for fp16:
+ // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ // smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ if constexpr (!use_inst_buffer) {
+ compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ } else {
+ if constexpr (!use_pv_fp16_accu) {
+ compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ } else {
+ compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ }
+ }
+
+ __syncthreads();
+ // load V
+ // for fp16:
+ // load_global_to_share stride_seq_v
+ load_fp8_V_global_to_share<global_to_shared_line_lanes_V, global_to_shared_copy_lines_per_warp_V, V_smem_iters_row, V_smem_iters_col, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, CTA_K>(
+ &V_lane_base_ptr, V_smem_offset_load, stride_d_v, smem_V);
+ cp_async::commit_group();
+ K_load_idx_lane_base += CTA_K;
+ }
+
+ // last iter, apply causal mask and out of bound mask
+ {
+ // ensure K is ready
+ cp_async::wait_group<1>();
+ __syncthreads();
+
+ // compute QK^T
+ if constexpr (num_tiles_qk_inner == 1) {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_K, RS, RQ, K_smem_offset_mma);
+ } else {
+ compute_int_qk<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_qk_inner, swizzle_mode_QK, QK_SMEM_STRIDE / PACK_SIZE_QK, DTypeQK>(
+ smem_Q, smem_K, RS, Q_smem_offset_mma, K_smem_offset_mma);
+ }
+
+ float RS_f32[num_tiles_q][num_tiles_k][8];
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ RS_f32[fq][fk][k] = __int2float_rz(RS[fq][fk][k]) * dequant_scale;
+ }
+ }
+ }
+
+ if constexpr (mask_mode == MaskMode::kCausal) {
+ apply_causal_mask<num_tiles_q, num_tiles_k>(Q_idx_lane_base, K_idx_lane_base, RS_f32);
+ }
+ apply_out_of_bound_mask<num_tiles_q, num_tiles_k>(K_idx_lane_base, RS_f32, kv_len);
+ K_idx_lane_base += CTA_K;
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, false, true, false>(RS_f32, RO, m, d, original_sm_scale);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ update_mdo<num_tiles_q, num_tiles_k, num_tiles_v, true, true, false>(RS_f32, RO, m, d, original_sm_scale);
+ }
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kCudaCore) {
+ accumulate_d<num_tiles_q, num_tiles_k, ComputeUnit::kCudaCore>(RS_f32, d);
+ }
+
+ uint32_t RS_f8[num_tiles_q][num_tiles_k / 2][4];
+ RS_32_to_8<num_tiles_q, num_tiles_k>(RS_f32, RS_f8);
+
+ if constexpr (DenominatorAccumUnit == ComputeUnit::kTensorCore) {
+ accumulate_d_f8<num_tiles_q, num_tiles_k>(RS_f8, d);
+ }
+
+ // ensure V is ready
+ cp_async::wait_group<0>();
+ __syncthreads();
+
+ // for fp16:
+ // compute_fp16_sv_permuted<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V, 4>(
+ // smem_V, RS_f16, RO, d, V_smem_offset_mma);
+ if constexpr (!use_inst_buffer) {
+ compute_fp8_sv<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ } else {
+ if constexpr (!use_pv_fp16_accu) {
+ compute_fp8_sv_inst_buf<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ } else {
+ compute_fp8_sv_inst_buf_fp16_accu<num_warps_q, num_warps_k, num_tiles_q, num_tiles_k, num_tiles_v, swizzle_mode_V, V_SMEM_STRIDE / PACK_SIZE_V>(
+ smem_V, RS_f8, RO, d);
+ }
+ }
+
+ __syncthreads();
+ }
+
+ // TODO: thread block sync mdo state for num_warps_k > 0. Then only one thread block needs to do the final saving.
+
+ normalize_d<num_tiles_q, num_tiles_v, ComputeUnit::kCudaCore>(RO, m, d);
+
+ // ! here we just implement the case for fp32 acumulation
+ if constexpr (fuse_v_scale) {
+ float v_scale[4];
+ float* V_scale_base_ptr = V_scale + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ ((float2*)v_scale)[0] = *((float2*)(V_scale_base_ptr + fv * 16));
+ ((float2*)v_scale)[1] = *((float2*)(V_scale_base_ptr + fv * 16 + 8));
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ RO[fq][fv][0] *= v_scale[0];
+ RO[fq][fv][1] *= v_scale[1];
+ RO[fq][fv][2] *= v_scale[0];
+ RO[fq][fv][3] *= v_scale[1];
+ RO[fq][fv][4] *= v_scale[2];
+ RO[fq][fv][5] *= v_scale[3];
+ RO[fq][fv][6] *= v_scale[2];
+ RO[fq][fv][7] *= v_scale[3];
+ }
+ }
+ }
+
+ if constexpr (fuse_v_mean) {
+ float v_mean[4];
+ float* V_mean_base_ptr = V_mean + batch_id * (num_qo_heads / num_kv_groups) * head_dim + (head_id / num_kv_groups) * head_dim + (lane_id % 4) * 2;
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ ((float2*)v_mean)[0] = *((float2*)(V_mean_base_ptr + fv * 16));
+ ((float2*)v_mean)[1] = *((float2*)(V_mean_base_ptr + fv * 16 + 8));
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ RO[fq][fv][0] += v_mean[0];
+ RO[fq][fv][1] += v_mean[1];
+ RO[fq][fv][2] += v_mean[0];
+ RO[fq][fv][3] += v_mean[1];
+ RO[fq][fv][4] += v_mean[2];
+ RO[fq][fv][5] += v_mean[3];
+ RO[fq][fv][6] += v_mean[2];
+ RO[fq][fv][7] += v_mean[3];
+ }
+ }
+ }
+
+ // save the result to shared memory
+ uint32_t smem_O_row_base = get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / 4;
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ uint32_t offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O));
+
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ // convert RO to half
+ uint32_t RO_f16[4];
+#pragma unroll
+ for (uint32_t k = 0; k < 4; k++) {
+ if constexpr (std::is_same<DTypeOut, half>::value) {
+ ((half2*)RO_f16)[k] = __float22half2_rn(((float2*)RO[fq][fv])[k]);
+ } else {
+ ((nv_bfloat162*)RO_f16)[k] = __float22bfloat162_rn(((float2*)RO[fq][fv])[k]);
+ }
+ }
+
+ ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[0];
+ ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[1];
+
+ offset_O = smem_O.get_permuted_offset(smem_O_row_base + fq * MMA_QK_M, fv * (MMA_SV_N / PACK_SIZE_O) + 1);
+ ((uint32_t*)(smem_O.base + offset_O))[lane_id % 4] = RO_f16[2];
+ ((uint32_t*)(smem_O.base + offset_O + 8 * (O_SMEM_STRIDE / PACK_SIZE_O)))[lane_id % 4] = RO_f16[3];
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ // TODO: not implement
+ }
+ }
+ }
+
+ // ! do we need to sync here?
+ __syncwarp();
+
+ // shared memory to global memory
+ DTypeOut* O_lane_ptr = O + batch_id * stride_bz_o + head_id * stride_h_o + (bx * CTA_Q + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>() + lane_id / global_to_shared_line_lanes_O) * stride_seq_o + lane_id % global_to_shared_line_lanes_O * PACK_SIZE_O;
+ uint32_t offset_O = smem_O.get_permuted_offset(get_warp_idx_q<num_warps_q, num_warps_k>() * WARP_Q + lane_id / global_to_shared_line_lanes_O, lane_id % global_to_shared_line_lanes_O);
+ uint32_t O_load_idx_lane_base = bx * CTA_Q + CTA_Q / num_warps * warp_id + lane_id / global_to_shared_line_lanes_O;
+
+#pragma unroll
+ for (uint32_t i = 0; i < O_smem_iters_col; i++) {
+#pragma unroll
+ for (uint32_t j = 0; j < O_smem_iters_row; j++) {
+ if (O_load_idx_lane_base < qo_len) {
+ smem_O.store_128b(offset_O, O_lane_ptr);
+ }
+ O_lane_ptr += (global_to_shared_line_lanes_O * PACK_SIZE_O);
+ offset_O = smem_O.advance_offset_by_column<global_to_shared_line_lanes_O>(offset_O);
+ }
+
+ offset_O = smem_O.advance_offset_by_row<global_to_shared_copy_lines_per_warp_O>(offset_O - (O_smem_iters_row * global_to_shared_line_lanes_O));
+ O_lane_ptr += ((global_to_shared_copy_lines_per_warp_O * stride_seq_o) - (O_smem_iters_row * global_to_shared_line_lanes_O * PACK_SIZE_O));
+ O_load_idx_lane_base += global_to_shared_copy_lines_per_warp_O;
+ }
+
+ if constexpr (return_lse) {
+ // ! this only works for num_tiles_q = 2
+ uint32_t lse_idx = bx * CTA_Q + lane_id / 4 + 8 * (lane_id % 4) + WARP_Q * get_warp_idx_q<num_warps_q, num_warps_k>();
+ float* lse_lane_ptr = Lse + batch_id * (qo_len * num_qo_heads) + head_id * qo_len + lse_idx;
+ uint32_t fq = (lane_id % 4) / 2;
+ uint32_t k = (lane_id % 4) % 2;
+
+ if (lse_idx < qo_len) {
+ lse_lane_ptr[0] = (math::ptx_log2(d[fq][k]) + m[fq][k]);
+ }
+ }
+#else
+ RUNTIME_ASSERT("FP8 SageAttention requires SM89 or newer");
+#endif
+ }
+
+#undef PACK_SIZE_QK
+#undef PACK_SIZE_V
+#undef PACK_SIZE_O
+#undef MMA_QK_M
+#undef MMA_QK_N
+#undef MMA_QK_K
+#undef MMA_SV_M
+#undef MMA_SV_N
+#undef MMA_SV_K
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/attn_utils.cuh b/src/core/cuda/sageattention/attn_utils.cuh
new file mode 100644
index 0000000..45157b5
--- /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 <cuda_fp16.h>
+#include <cuda_pipeline_primitives.h>
+
+#include "cp_async.cuh"
+#include "math.cuh"
+#include "mma.cuh"
+#include "numeric_conversion.cuh"
+#include "permuted_smem.cuh"
+
+namespace sageattention {
+
+#define WARP_SIZE 32
+
+#define S_FP8_OFFSET 8.807f
+#define S_FP8_OFFSET_EXP 6680.8477f
+#define S_FP8_OFFSET_EXP_INV 0.0022326917f
+
+#define SAGE_DIV_CEIL(M, N) (((M) + (N) - 1) / (N))
+
+ enum class MaskMode {
+ kNone = 0,
+ kCausal = 1,
+ };
+
+ enum class DataType {
+ kHalf,
+ kInt8,
+ kInt4,
+ kE4M3,
+ kE5M2,
+ };
+
+ enum class QuantGranularity {
+ kPerTensor = 0,
+ kPerBlock = 1,
+ kPerWarp = 2,
+ kPerThread = 3,
+ };
+
+ enum class ComputeUnit {
+ kTensorCore,
+ kCudaCore,
+ };
+
+ __device__ __forceinline__ uint32_t get_warp_id() {
+ return threadIdx.y;
+ }
+
+ __device__ __forceinline__ uint32_t get_lane_id() {
+ return threadIdx.x;
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k>
+ __device__ __forceinline__ uint32_t get_warp_idx_q() {
+ return get_warp_id() / num_warps_k;
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k>
+ __device__ __forceinline__ uint32_t get_warp_idx_k() {
+ return get_warp_id() % num_warps_k;
+ }
+
+ template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
+ __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
+ static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+ static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
+
+ constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
+
+#pragma unroll
+ for (uint32_t i = 0; i < smem_iters_col; i++) {
+#pragma unroll
+ for (uint32_t j = 0; j < smem_iters_row; j++) {
+ smem.load_128b_async(smem_offset, *lane_ptr);
+ *lane_ptr += (global_to_shared_line_lanes * pack_size);
+ smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
+ }
+
+ smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
+ *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size));
+ }
+ smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
+ *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+ }
+
+ // with predicate
+ template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA, typename T>
+ __device__ __forceinline__ void load_global_to_share(T** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem, uint32_t base_idx, uint32_t max_len) {
+ static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+ static_assert(std::is_same<T, half>::value || std::is_same<T, int8_t>::value);
+
+ constexpr uint32_t pack_size = std::is_same<T, half>::value ? 8 : 16;
+
+#pragma unroll
+ for (uint32_t i = 0; i < smem_iters_col; i++) {
+#pragma unroll
+ for (uint32_t j = 0; j < smem_iters_row; j++) {
+ // Masked probabilities still propagate NaNs from stale V lanes through 0 * NaN.
+ smem.load_128b_async<std::is_same<T, half>::value ? cp_async::SharedMemFillMode::kFillZero : cp_async::SharedMemFillMode::kNoFill>(smem_offset, *lane_ptr, base_idx < max_len);
+ *lane_ptr += (global_to_shared_line_lanes * pack_size);
+ smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
+ }
+
+ smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
+ *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size));
+ base_idx += global_to_shared_copy_lines_per_warp_per_iter;
+ }
+ smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
+ *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+ }
+
+ template <uint32_t global_to_shared_line_lanes, uint32_t global_to_shared_copy_lines_per_warp_per_iter, uint32_t smem_iters_row, uint32_t smem_iters_col, SwizzleMode swizzle_mode, uint32_t stride, uint32_t CTA>
+ __device__ __forceinline__ void load_fp8_V_global_to_share(int8_t** lane_ptr, uint32_t& smem_offset, const uint32_t& gmem_stride, const smem_t<swizzle_mode, stride>& smem) {
+ static_assert(global_to_shared_copy_lines_per_warp_per_iter * global_to_shared_line_lanes == WARP_SIZE);
+ constexpr uint32_t pack_size_fp8 = 16;
+
+#pragma unroll
+ for (uint32_t i = 0; i < smem_iters_col; i++) {
+#pragma unroll
+ for (uint32_t j = 0; j < smem_iters_row; j++) {
+ smem.load_128b_async(smem_offset, *lane_ptr);
+ *lane_ptr += (global_to_shared_line_lanes * pack_size_fp8);
+ smem_offset = smem.advance_offset_by_column<global_to_shared_line_lanes>(smem_offset);
+ }
+
+ smem_offset = smem.advance_offset_by_row<global_to_shared_copy_lines_per_warp_per_iter>(smem_offset - (smem_iters_row * global_to_shared_line_lanes));
+ *lane_ptr += ((global_to_shared_copy_lines_per_warp_per_iter * gmem_stride) - (smem_iters_row * global_to_shared_line_lanes * pack_size_fp8));
+ }
+ smem_offset -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter * stride);
+ // for QK: *lane_ptr += (CTA - smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+ *lane_ptr += CTA; // ! prevent underflow
+ *lane_ptr -= (smem_iters_col * global_to_shared_copy_lines_per_warp_per_iter) * gmem_stride;
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
+ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_Q, const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t& offset_Q, uint32_t& offset_K) {
+ static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
+
+ uint32_t RQ[num_tiles_q][4];
+ uint32_t RK[4];
+
+ // the first iteration, mma mode is kInit
+#pragma unroll
+ for (uint32_t iter = 0; iter < 1; iter++) {
+ // load RQ
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
+ offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
+ }
+ // ! using permutation invariance
+ offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
+
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ // load RK
+ smem_K.ldmatrix_m8n8x4(offset_K, RK);
+ offset_K = smem_K.advance_offset_by_row<16>(offset_K);
+
+ // mma
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (DTypeQK == DataType::kInt8) {
+ mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+ } else if constexpr (DTypeQK == DataType::kInt4) {
+ mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+ }
+ }
+ }
+ offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter);
+ }
+
+ // following iteration, mma mode is kInplace
+#pragma unroll
+ for (uint32_t iter = 1; iter < num_tiles_qk_inner; iter++) {
+ // load RQ
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ smem_Q.ldmatrix_m8n8x4(offset_Q, RQ[fq]);
+ offset_Q = smem_Q.advance_offset_by_row<16>(offset_Q);
+ }
+ offset_Q = smem_Q.advance_offset_by_column<2>(offset_Q - (num_tiles_q * 16 * stride), iter);
+
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ // load RK
+ smem_K.ldmatrix_m8n8x4(offset_K, RK);
+ offset_K = smem_K.advance_offset_by_row<16>(offset_K);
+
+ // mma
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (DTypeQK == DataType::kInt8) {
+ mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
+ } else if constexpr (DTypeQK == DataType::kInt4) {
+ mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInplaceUpdate>(RS[fq][fk], RQ[fq], RK);
+ }
+ }
+ }
+ offset_K = smem_K.advance_offset_by_column<2>(offset_K - (num_tiles_k * 16 * stride), iter);
+ }
+
+ offset_Q -= (2 * num_tiles_qk_inner);
+ offset_K -= (2 * num_tiles_qk_inner);
+ }
+
+ // for case when num_tiles_qk_inner = 1
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_qk_inner, SwizzleMode swizzle_mode, uint32_t stride, DataType DTypeQK>
+ __device__ __forceinline__ void compute_int_qk(const smem_t<swizzle_mode, stride>& smem_K, int32_t RS[][num_tiles_k][8], uint32_t RQ[][4], uint32_t offset_K) {
+ static_assert(DTypeQK == DataType::kInt8 || DTypeQK == DataType::kInt4);
+ static_assert(num_tiles_qk_inner == 1);
+
+ uint32_t RK[4];
+
+ // mma mode is kInit
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ // load RK
+ smem_K.ldmatrix_m8n8x4(offset_K, RK);
+ offset_K = smem_K.advance_offset_by_row<16>(offset_K);
+
+ // mma
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (DTypeQK == DataType::kInt8) {
+ mma::mma_sync_m16n16k32_row_col_s8s8s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+ } else if constexpr (DTypeQK == DataType::kInt4) {
+ mma::mma_sync_m16n16k64_row_col_s4s4s32<mma::MMAMode::kInit>(RS[fq][fk], RQ[fq], RK);
+ }
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
+ __device__ __forceinline__ void apply_causal_mask(const uint32_t& Q_idx_lane_base, const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8]) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ const uint32_t q_idx = Q_idx_lane_base + fq * 16 + 8 * ((k % 4) / 2);
+ const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
+ const bool out_of_boundary = (kv_idx > q_idx);
+
+ if constexpr (std::is_same<DTypeQKAccum, float>::value) {
+ RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
+ } else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
+ RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
+ }
+ }
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename DTypeQKAccum>
+ __device__ __forceinline__ void apply_out_of_bound_mask(const uint32_t& K_idx_lane_base, DTypeQKAccum RS[][num_tiles_k][8], const uint32_t& kv_len) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ const uint32_t kv_idx = K_idx_lane_base + fk * 16 + 8 * (k / 4) + k % 2;
+ const bool out_of_boundary = (kv_idx >= kv_len);
+
+ if constexpr (std::is_same<DTypeQKAccum, float>::value) {
+ RS[fq][fk][k] = (out_of_boundary ? -5000000.0f : RS[fq][fk][k]);
+ } else if constexpr (std::is_same<DTypeQKAccum, half>::value) {
+ RS[fq][fk][k] = (out_of_boundary ? __float2half_rn(-50000.0f) : RS[fq][fk][k]);
+ }
+ }
+ }
+ }
+ }
+
+ // for DTypeQKAccum float
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, bool use_half_o_scale, bool exp_offset, bool fuse_scale = false, typename DTypeSVAccum>
+ __device__ __forceinline__ void update_mdo(float RS[][num_tiles_k][8], DTypeSVAccum RO[][num_tiles_v][8], float m[][2], float d[][2], const float& sm_scale) {
+ static_assert(std::is_same<DTypeSVAccum, half>::value || (!use_half_o_scale));
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 2; k++) {
+ // assign the smallest value possible
+ float m_prev = m[fq][k];
+ float m_temp = -5000000.0f;
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ float m_local = max(max(RS[fq][fk][k * 2 + 0], RS[fq][fk][k * 2 + 1]),
+ max(RS[fq][fk][k * 2 + 4], RS[fq][fk][k * 2 + 5]));
+ m_temp = max(m_temp, m_local);
+ }
+
+ if constexpr (!fuse_scale) {
+ if constexpr (exp_offset) {
+ m_temp = fmaf(m_temp, sm_scale, -S_FP8_OFFSET);
+ } else {
+ m_temp *= sm_scale;
+ }
+ } else if constexpr (exp_offset) {
+ m_temp += (-S_FP8_OFFSET);
+ }
+
+ // exchange element with the 4 threads in the row
+ m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x1)); // 0 exchange with 1, 2 exchange with 3
+ m_temp = max(m_temp, __shfl_xor_sync(0xffffffff, m_temp, 0x2)); // 0 exchange with 2, 1 exchange with 3
+
+ m[fq][k] = max(m[fq][k], m_temp);
+
+ float o_scale = math::ptx_exp2(m_prev - m[fq][k]);
+
+ // update denominator
+ d[fq][k] *= o_scale;
+
+ half2 o_scale2;
+ if constexpr (use_half_o_scale) {
+ o_scale2 = __floats2half2_rn(o_scale, o_scale);
+ }
+
+ // update RO
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ RO[fq][fv][k * 2 + 0] *= o_scale;
+ RO[fq][fv][k * 2 + 1] *= o_scale;
+ RO[fq][fv][k * 2 + 4] *= o_scale;
+ RO[fq][fv][k * 2 + 5] *= o_scale;
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ if constexpr (use_half_o_scale) {
+ ((half2*)RO[fq][fv])[k] = __hmul2(((half2*)RO[fq][fv])[k], o_scale2);
+ ((half2*)RO[fq][fv])[k + 2] = __hmul2(((half2*)RO[fq][fv])[k + 2], o_scale2);
+ } else {
+ RO[fq][fv][k * 2 + 0] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 0]) * o_scale);
+ RO[fq][fv][k * 2 + 1] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 1]) * o_scale);
+ RO[fq][fv][k * 2 + 4] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 4]) * o_scale);
+ RO[fq][fv][k * 2 + 5] = __float2half_rn(__half2float(RO[fq][fv][k * 2 + 5]) * o_scale);
+ }
+ }
+ }
+
+ // raise RS to exponent
+ float negative_m = -m[fq][k];
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ if constexpr (fuse_scale) {
+ RS[fq][fk][k * 2 + 0] = math::ptx_exp2(RS[fq][fk][k * 2 + 0] + negative_m);
+ RS[fq][fk][k * 2 + 1] = math::ptx_exp2(RS[fq][fk][k * 2 + 1] + negative_m);
+ RS[fq][fk][k * 2 + 4] = math::ptx_exp2(RS[fq][fk][k * 2 + 4] + negative_m);
+ RS[fq][fk][k * 2 + 5] = math::ptx_exp2(RS[fq][fk][k * 2 + 5] + negative_m);
+ } else {
+ RS[fq][fk][k * 2 + 0] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 0], sm_scale, negative_m));
+ RS[fq][fk][k * 2 + 1] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 1], sm_scale, negative_m));
+ RS[fq][fk][k * 2 + 4] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 4], sm_scale, negative_m));
+ RS[fq][fk][k * 2 + 5] = math::ptx_exp2(fmaf(RS[fq][fk][k * 2 + 5], sm_scale, negative_m));
+ }
+ }
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k, typename T>
+ __device__ __forceinline__ void RS_32_to_16(T RS[][num_tiles_k][8], uint32_t RS_16[][num_tiles_k][4]) {
+ static_assert(sizeof(T) == 4);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ ((half2*)RS_16[fq][fk])[0] = __float22half2_rn(((float2*)RS[fq][fk])[0]);
+ ((half2*)RS_16[fq][fk])[1] = __float22half2_rn(((float2*)RS[fq][fk])[1]);
+ ((half2*)RS_16[fq][fk])[2] = __float22half2_rn(((float2*)RS[fq][fk])[2]);
+ ((half2*)RS_16[fq][fk])[3] = __float22half2_rn(((float2*)RS[fq][fk])[3]);
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+ __device__ __forceinline__ void RS_32_to_8(float RS[][num_tiles_k][8], uint32_t RS_8[][num_tiles_k / 2][4]) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+ floatx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 4);
+ floatx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 2, RS[fq][fk * 2 + 0] + 6);
+ floatx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 4);
+ floatx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 2, RS[fq][fk * 2 + 1] + 6);
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+ __device__ __forceinline__ void RS_16_to_8(uint32_t RS[][num_tiles_k][4], uint32_t RS_8[][num_tiles_k / 2][4]) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+ halfx4_to_e4m3x4(RS_8[fq][fk], RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2);
+ halfx4_to_e4m3x4(RS_8[fq][fk] + 1, RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3);
+ halfx4_to_e4m3x4(RS_8[fq][fk] + 2, RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2);
+ halfx4_to_e4m3x4(RS_8[fq][fk] + 3, RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3);
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+ __device__ __forceinline__ void RS_8_to_16(uint32_t RS_8[][num_tiles_k / 2][4], uint32_t RS[][num_tiles_k][4]) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+ e4m3x4_to_halfx4(RS[fq][fk * 2 + 0], RS[fq][fk * 2 + 0] + 2, RS_8[fq][fk]);
+ e4m3x4_to_halfx4(RS[fq][fk * 2 + 0] + 1, RS[fq][fk * 2 + 0] + 3, RS_8[fq][fk] + 1);
+ e4m3x4_to_halfx4(RS[fq][fk * 2 + 1], RS[fq][fk * 2 + 1] + 2, RS_8[fq][fk] + 2);
+ e4m3x4_to_halfx4(RS[fq][fk * 2 + 1] + 1, RS[fq][fk * 2 + 1] + 3, RS_8[fq][fk] + 3);
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k, ComputeUnit compute_unit = ComputeUnit::kTensorCore, typename T>
+ __device__ __forceinline__ void accumulate_d(T RS[][num_tiles_k][(compute_unit == ComputeUnit::kTensorCore) ? 4 : 8], float d[][2]) {
+ // for compute unit cuda core, RS is float
+ // for compute unit tensor core, RS is packed half
+ static_assert((std::is_same<T, float>::value && compute_unit == ComputeUnit::kCudaCore) ||
+ (std::is_same<T, uint32_t>::value && compute_unit == ComputeUnit::kTensorCore));
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+ if constexpr (compute_unit == ComputeUnit::kTensorCore) {
+ // full accumulate with tensor core
+ mma::rowsum_f16f16f32(d[fq], (uint32_t*)(RS[fq][fk]));
+ } else if constexpr (compute_unit == ComputeUnit::kCudaCore) {
+ // partial accumulate with cuda core
+ d[fq][0] += RS[fq][fk][0] + RS[fq][fk][1] + RS[fq][fk][4] + RS[fq][fk][5];
+ d[fq][1] += RS[fq][fk][2] + RS[fq][fk][3] + RS[fq][fk][6] + RS[fq][fk][7];
+ }
+ }
+ }
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_k>
+ __device__ __forceinline__ void accumulate_d_f8(uint32_t RS[][num_tiles_k / 2][4], float d[][2]) {
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+ mma::rowsum_f8f8f32(d[fq], RS[fq][fk]);
+ }
+ }
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+ __device__ __forceinline__ void compute_fp16_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f16[][num_tiles_k][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+ uint32_t smem_V_row_base = get_warp_idx_k<num_warps_q, num_warps_k>() * (num_tiles_k * 16) + get_lane_id() % 16;
+ uint32_t smem_V_col_base = get_lane_id() / 16;
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fk * 16, smem_V_col_base + fv * 2);
+ smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], RS_f16[fq][fk], RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], RS_f16[fq][fk], RV);
+ }
+ }
+ }
+ }
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
+ __device__ __forceinline__ void compute_fp16_sv_permuted(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
+ static_assert(sizeof(T) == 4);
+
+ // ! be sure you know what you are doing
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ mma::mma_sync_m16n16k16_row_col_f16f16f32(RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ mma::mma_sync_m16n16k16_row_col_f16f16f16((uint32_t*)RO[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+ }
+ }
+
+ offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
+ }
+
+ // make offset_V their original value
+ offset_V -= (16 * num_tiles_k * stride);
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, uint32_t RS_width = 4, typename T, typename DTypeSVAccum>
+ __device__ __forceinline__ void compute_fp16_sv_permuted_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, T RS_f16[][num_tiles_k][RS_width], DTypeSVAccum RO[][num_tiles_v][8], float d[][2], uint32_t& offset_V) {
+ static_assert(sizeof(T) == 4);
+ static_assert(std::is_same<DTypeSVAccum, float>::value);
+
+ uint32_t RO_inst_buf[num_tiles_q][num_tiles_v][4];
+
+ // ! be sure you know what you are doing
+#pragma unroll
+ for (uint32_t fk = 0; fk < 1; fk++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ {
+ mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInit>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+ }
+ }
+
+ offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
+ }
+
+#pragma unroll
+ for (uint32_t fk = 1; fk < num_tiles_k; fk++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ smem_V.ldmatrix_m8n8x4_trans(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ {
+ mma::mma_sync_m16n16k16_row_col_f16f16f16<mma::MMAMode::kInplaceUpdate>((uint32_t*)RO_inst_buf[fq][fv], (uint32_t*)(RS_f16[fq][fk]), RV);
+ }
+ }
+
+ offset_V = smem_V.advance_offset_by_column<2>(offset_V, fv);
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V - (2 * num_tiles_v));
+ }
+
+ // accumulate into RO
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ RO[fq][fv][0] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].x);
+ RO[fq][fv][1] += __half2float(((half2*)RO_inst_buf[fq][fv])[0].y);
+ RO[fq][fv][2] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].x);
+ RO[fq][fv][3] += __half2float(((half2*)RO_inst_buf[fq][fv])[1].y);
+ RO[fq][fv][4] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].x);
+ RO[fq][fv][5] += __half2float(((half2*)RO_inst_buf[fq][fv])[2].y);
+ RO[fq][fv][6] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].x);
+ RO[fq][fv][7] += __half2float(((half2*)RO_inst_buf[fq][fv])[3].y);
+ }
+ }
+
+ // make offset_V their original value
+ offset_V -= (16 * num_tiles_k * stride);
+ }
+
+ template <uint32_t num_tiles_q, uint32_t num_tiles_v,
+ ComputeUnit compute_unit = ComputeUnit::kTensorCore, // compute unit for accumulate_d
+ typename DTypeQKAccum,
+ typename DTypeSVAccum>
+ __device__ __forceinline__ void normalize_d(DTypeSVAccum RO[][num_tiles_v][8], DTypeQKAccum m[][2], float d[][2]) {
+ if constexpr (compute_unit == ComputeUnit::kCudaCore) {
+ // accumulate_d performs partial accumulation with cuda core
+ // aggregate d
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 2; k++) {
+ d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x1); // sum 0 and 1, 2 and 3
+ d[fq][k] += __shfl_xor_sync(0xffffffff, d[fq][k], 0x2); // sum 0 and 2, 1 and 3
+ }
+ }
+ }
+
+ // divide O by d
+ float d_rcp[num_tiles_q][2];
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 2; k++) {
+ // TODO: check m to prevent nan
+ d_rcp[fq][k] = math::ptx_rcp(d[fq][k]);
+ }
+ }
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+#pragma unroll
+ for (uint32_t k = 0; k < 8; k++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ RO[fq][fv][k] *= d_rcp[fq][(k % 4) / 2];
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ RO[fq][fv][k] = __float2half_rn(__half2float(RO[fq][fv][k]) * d_rcp[fq][(k % 4) / 2]);
+ }
+ }
+ }
+ }
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+ __device__ __forceinline__ void compute_fp8_sv(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+ uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
+ // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
+ uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
+#pragma unroll
+ for (uint32_t fk = 0; fk < num_tiles_k / 2; fk++) {
+ uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+ smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ mma::mma_sync_m16n16k32_row_col_f8f8f32(RO[fq][fv], RS_f8[fq][fk], RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ // ! Not Implemented
+ }
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+ }
+ }
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+ __device__ __forceinline__ void compute_fp8_sv_inst_buf(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+ uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
+ // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
+ uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
+
+ float RO_inst_buf[num_tiles_q][num_tiles_v][8];
+
+#pragma unroll
+ for (uint32_t fk = 0; fk < 1; fk++) {
+ uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+ smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ // ! Not Implemented
+ }
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+ }
+ }
+
+#pragma unroll
+ for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
+ uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+ smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ // ! Not Implemented
+ }
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+ }
+ }
+
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
+ RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
+ RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
+ RO[fq][fv][3] += RO_inst_buf[fq][fv][3];
+ RO[fq][fv][4] += RO_inst_buf[fq][fv][4];
+ RO[fq][fv][5] += RO_inst_buf[fq][fv][5];
+ RO[fq][fv][6] += RO_inst_buf[fq][fv][6];
+ RO[fq][fv][7] += RO_inst_buf[fq][fv][7];
+ }
+ }
+ }
+
+ template <uint32_t num_warps_q, uint32_t num_warps_k, uint32_t num_tiles_q, uint32_t num_tiles_k, uint32_t num_tiles_v, SwizzleMode swizzle_mode, uint32_t stride, typename DTypeSVAccum>
+ __device__ __forceinline__ void compute_fp8_sv_inst_buf_fp16_accu(const smem_t<swizzle_mode, stride>& smem_V, uint32_t RS_f8[][num_tiles_k / 2][4], DTypeSVAccum RO[][num_tiles_v][8], float d[][2]) {
+ uint32_t smem_V_row_base = get_lane_id() % 8 + (get_lane_id() / 16) * 8;
+ // uint32_t smem_V_col_base = get_warp_idx_k<num_warps_q, num_warps_k>() * ((16 * num_tiles_k) / 16) + (get_lane_id() / 8) % 2;
+ uint32_t smem_V_col_base = (get_lane_id() / 8) % 2;
+
+ uint32_t RO_int32[num_tiles_q][num_tiles_v][4];
+
+#pragma unroll
+ for (uint32_t fk = 0; fk < 1; fk++) {
+ uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+ smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ // mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInit>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+ mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInit>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ // ! Not Implemented
+ }
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+ }
+ }
+
+#pragma unroll
+ for (uint32_t fk = 1; fk < num_tiles_k / 2; fk++) {
+ uint32_t offset_V = smem_V.get_permuted_offset(smem_V_row_base, smem_V_col_base + fk * 2);
+#pragma unroll
+ for (uint32_t fv = 0; fv < num_tiles_v; fv++) {
+ // load RV
+ uint32_t RV[4];
+ // uint32_t offset_V = (smem_V).get_permuted_offset(smem_V_row_base + fv * 16, smem_V_col_base + fk * 2);
+ smem_V.ldmatrix_m8n8x4(offset_V, RV);
+#pragma unroll
+ for (uint32_t fq = 0; fq < num_tiles_q; fq++) {
+ if constexpr (std::is_same<DTypeSVAccum, float>::value) {
+ // mma::mma_sync_m16n16k32_row_col_f8f8f32<mma::MMAMode::kInplaceUpdate>(RO_inst_buf[fq][fv], RS_f8[fq][fk], RV);
+ mma::mma_sync_m16n16k32_row_col_f8f8f16<mma::MMAMode::kInplaceUpdate>(RO_int32[fq][fv], RS_f8[fq][fk], RV);
+ } else if constexpr (std::is_same<DTypeSVAccum, half>::value) {
+ // ! Not Implemented
+ }
+ }
+ offset_V = smem_V.advance_offset_by_row<16>(offset_V);
+ }
+ }
+ float RO_tmp_float[2];
+#pragma unroll
+ for (int i = 0; i < num_tiles_q; i++) {
+#pragma unroll
+ for (int j = 0; j < num_tiles_v; j++) {
+#pragma unroll
+ for (int k = 0; k < 4; k++) {
+ unpack_half2_from_uint32_to_float(RO_tmp_float, RO_int32[i][j][k]);
+ RO[i][j][k * 2 + 0] += RO_tmp_float[0];
+ RO[i][j][k * 2 + 1] += RO_tmp_float[1];
+ }
+ }
+ }
+
+ // #pragma unroll
+ // for (uint32_t fq = 0; fq < num_tiles_q; fq++)
+ // {
+ // #pragma unroll
+ // for (uint32_t fv = 0; fv < num_tiles_v; fv++)
+ // {
+ // RO[fq][fv][0] += RO_inst_buf[fq][fv][0];
+ // RO[fq][fv][1] += RO_inst_buf[fq][fv][1];
+ // RO[fq][fv][2] += RO_inst_buf[fq][fv][2];
+ // RO[fq][fv][3] += RO_inst_buf[fq][fv][3];
+ // RO[fq][fv][4] += RO_inst_buf[fq][fv][4];
+ // RO[fq][fv][5] += RO_inst_buf[fq][fv][5];
+ // RO[fq][fv][6] += RO_inst_buf[fq][fv][6];
+ // RO[fq][fv][7] += RO_inst_buf[fq][fv][7];
+ // }
+ // }
+ }
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/cp_async.cuh b/src/core/cuda/sageattention/cp_async.cuh
new file mode 100644
index 0000000..98ceca5
--- /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 <cuda_runtime.h>
+
+namespace sageattention {
+
+ namespace cp_async {
+
+ enum class SharedMemFillMode {
+ kFillZero, // Fill zero to shared memory when predicate is false
+ kNoFill // Do not fill zero to shared memory when predicate is false
+ };
+
+ enum class PrefetchMode {
+ kNoPrefetch, // Do not fetch additional data from global memory to L2
+ kPrefetch // Fetch additional data from global memory to L2
+ };
+
+#if (__CUDACC_VER_MAJOR__ >= 11)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
+#define CP_ASYNC_ENABLED
+#endif
+#endif
+
+ /*!
+ * \brief Wrapper of PTX cp.async.commit_group instruction, commit all prior uncommitted
+ * cp.async instructions to a group
+ */
+ __device__ __forceinline__ void commit_group() {
+#ifdef CP_ASYNC_ENABLED
+ asm volatile("cp.async.commit_group;\n" ::);
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of PTX cp.async.wait_group instruction
+ * \tparam n Wait till most recent n groups are committed
+ */
+ template <size_t n>
+ __device__ __forceinline__ void wait_group() {
+#ifdef CP_ASYNC_ENABLED
+ asm volatile("cp.async.wait_group %0;\n" ::"n"(n));
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
+ * global memory to shared memory
+ * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
+ * \tparam T Data type
+ * \param smem_ptr Pointer to shared memory
+ * \param gmem_ptr Pointer to global memory
+ */
+ template <PrefetchMode prefetch_mode, typename T>
+ __device__ __forceinline__ void load_128b(T* smem_ptr, const T* gmem_ptr) {
+#ifdef CP_ASYNC_ENABLED
+ uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+ if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+ asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+ "l"(gmem_ptr), "n"(16), "r"(16));
+ } else {
+ asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+ "l"(gmem_ptr), "n"(16), "r"(16));
+ }
+#else
+ *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of PTX cp.async.cg.shared.global instruction, asynchronously copy data from
+ * global memory to shared memory with predicate.
+ * \tparam prefetch_mode Whether to fetch additional data from global memory to L2
+ * \tparam fill_mode Whether to fill zero to shared memory when predicate is false
+ * \tparam T Data type
+ * \param smem_ptr Pointer to shared memory
+ * \param gmem_ptr Pointer to global memory
+ * \param predicate Predicate value
+ * \note fill zero is slower than not fill zero
+ */
+ template <PrefetchMode prefetch_mode, SharedMemFillMode fill_mode, typename T>
+ __device__ __forceinline__ void pred_load_128b(T* smem_ptr, const T* gmem_ptr, bool predicate) {
+#ifdef CP_ASYNC_ENABLED
+ uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+ if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+ int src_in_bytes = predicate ? 16 : 0;
+ if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+ asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+ "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
+ } else {
+ asm volatile("cp.async.cg.shared.global [%0], [%1], %2, %3;\n" ::"r"(smem_int_ptr),
+ "l"(gmem_ptr), "n"(16), "r"(src_in_bytes));
+ }
+ } else {
+ if constexpr (prefetch_mode == PrefetchMode::kPrefetch) {
+ asm volatile(
+ "{\n"
+ " .reg .pred p;\n"
+ " setp.ne.b32 p, %0, 0;\n"
+ " @p cp.async.cg.shared.global.L2::128B [%1], [%2], %3;\n"
+ "}\n" ::"r"((int)predicate),
+ "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
+ } else {
+ asm volatile(
+ "{\n"
+ " .reg .pred p;\n"
+ " setp.ne.b32 p, %0, 0;\n"
+ " @p cp.async.cg.shared.global [%1], [%2], %3;\n"
+ "}\n" ::"r"((int)predicate),
+ "r"(smem_int_ptr), "l"(gmem_ptr), "n"(16));
+ }
+ }
+#else
+ if (predicate) {
+ *((uint4*)smem_ptr) = *((uint4*)gmem_ptr);
+ } else {
+ if constexpr (fill_mode == SharedMemFillMode::kFillZero) {
+ *((uint4*)smem_ptr) = make_uint4(0, 0, 0, 0);
+ }
+ }
+#endif
+ }
+
+ } // namespace cp_async
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/math.cuh b/src/core/cuda/sageattention/math.cuh
new file mode 100644
index 0000000..b92ac79
--- /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 <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+namespace sageattention {
+#ifndef USHORT_TYPE
+#define USHORT_TYPE
+ typedef unsigned short ushort;
+#endif
+
+ namespace math {
+
+ // log2(e)
+ constexpr float log2e = 1.44269504088896340736f;
+ constexpr float log2e_recp = 1.0f / log2e;
+
+ __forceinline__ __device__ half2 uint32_as_half2(uint32_t x) {
+ return *(half2*)&x;
+ }
+
+ __forceinline__ __device__ uint32_t half2_as_uint32(half2 x) {
+ return *(uint32_t*)&x;
+ }
+
+ /*!
+ * \brief Wrapper of PTX ex2.approx instruction, which computes 2^x
+ * \param x input
+ */
+ __forceinline__ __device__ float ptx_exp2(float x) {
+ float y;
+ asm volatile("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+ return y;
+ }
+
+ /*!
+ * \brief Wrapper of PTX lg2.approx instruction, which computes log2(x)
+ * \param x input
+ */
+ __forceinline__ __device__ float ptx_log2(float x) {
+ float y;
+ asm volatile("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+ return y;
+ }
+
+ /*!
+ * \brief Wrapper of PTX ex2.approx.f16x2 instruction, which computes 2^x
+ * \param x input
+ */
+ __forceinline__ __device__ half2 ptx_exp2(half2 x) {
+ uint32_t y_u32;
+ uint32_t x_u32 = half2_as_uint32(x);
+ asm volatile("ex2.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
+ return uint32_as_half2(y_u32);
+ }
+
+ /*!
+ * \brief Wrapper of PTX ex2.approx.f16 instruction, which computes 2^x
+ * \param x input
+ */
+ __forceinline__ __device__ half ptx_exp2(half x) {
+ ushort y_u16;
+ asm volatile("ex2.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
+ return __ushort_as_half(y_u16);
+ }
+
+ /*!
+ * \brief Wrapper of PTX rcp.approx instruction, which computes 1/x
+ * \param x input
+ */
+ __forceinline__ __device__ float ptx_rcp(float x) {
+ float y;
+ asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+ return y;
+ }
+
+ /*!
+ * \brief Wrapper of PTX shfl.sync.bfly instruction, which performs a butterfly shuffle
+ * between threads in a warp.
+ * \param x The value in the source lane
+ * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ delta]
+ */
+ __forceinline__ __device__ float shfl_xor_sync(float x, int lane_mask) {
+ float y;
+ asm volatile("shfl.sync.bfly.b32 %0, %1, %2, 0x1f, 0xffffffff;"
+ : "=f"(y)
+ : "f"(x), "r"(lane_mask));
+ return y;
+ }
+
+ /*!
+ * \brief Wrapper of PTX shfl.sync.bfly instruction on half2, which performs a butterfly
+ * shuffle between threads in a warp.
+ * \param x The value in the source lane
+ * \param lane_mask The mask to perform thread index xor with: y[i] <- x[i ^ lane_mask]
+ */
+ __forceinline__ __device__ half2 shfl_xor_sync(half2 x, int lane_mask) {
+ return __shfl_xor_sync(0xffffffff, x, lane_mask);
+ }
+
+ /*!
+ * \brief Wrapper of PTX rsqrt approximation instruction, which computes 1/sqrt(x)
+ * \param x input
+ */
+ __forceinline__ __device__ float rsqrt(float x) {
+ float y;
+ asm volatile("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x));
+ return y;
+ }
+
+ /*!
+ * \brief Wrapper of PTX tanh.approx.f32 instruction, which computes tanh(x)
+ * \param x input
+ */
+ __forceinline__ __device__ float tanh(float x) {
+ float y;
+ asm volatile("tanh.approx.f32 %0, %1;" : "=f"(y) : "f"(x));
+ return y;
+ }
+
+ /*!
+ * \brief Wrapper of PTX tanh.approx.f16x2 instruction, which computes tanh(x)
+ * \param x input
+ */
+ __forceinline__ __device__ half2 tanh(half2 x) {
+ uint32_t y_u32;
+ uint32_t x_u32 = half2_as_uint32(x);
+ asm volatile("tanh.approx.f16x2 %0, %1;" : "=r"(y_u32) : "r"(x_u32));
+ return uint32_as_half2(y_u32);
+ }
+
+ /*!
+ * \brief Wrapper of PTX tanh.approx.f16 instruction, which computes tanh(x)
+ * \param x input
+ */
+ __forceinline__ __device__ half tanh(half x) {
+ ushort y_u16;
+ asm volatile("tanh.approx.f16 %0, %1;" : "=h"(y_u16) : "h"(__half_as_ushort(x)));
+ return __ushort_as_half(y_u16);
+ }
+
+ } // namespace math
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/mma.cuh b/src/core/cuda/sageattention/mma.cuh
new file mode 100644
index 0000000..4de4828
--- /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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+namespace sageattention {
+
+ namespace mma {
+
+#if (__CUDACC_VER_MAJOR__ >= 11)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 800))
+#define MMA_F16F16F32_M16N8K16_ENABLED
+#define MMA_F16F16F16_M16N8K16_ENABLED
+#define MMA_S8S8S32_M16N8K32_ENABLED
+#define MMA_S4S4S32_M16N8K64_ENABLED
+#endif
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 750))
+#define MMA_F16F16F32_M16N8K8_ENABLED
+#define MMA_F16F16F16_M16N8K8_ENABLED
+#define LDMATRIX_M8N8X2_ENABLED
+#define LDMATRIX_M8N8X4_ENABLED
+#endif
+#endif
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
+#define MMA_F8F8F32_M16N8K16_ENABLED
+#endif
+#endif
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120800)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
+#define MMA_F8F8F16_M16N8K16_ENABLED
+#endif
+#endif
+
+#if defined(__CUDA_ARCH__)
+#define RUNTIME_ASSERT(x) __brkpt()
+#else
+
+#define RUNTIME_ASSERT(x) assert(0 && x)
+#endif
+
+ enum class MMAMode {
+ kInit = 0U,
+ kInplaceUpdate = 1U,
+ };
+
+ /*!
+ * \brief Wrapper of PTX ldmatrix m8n8.x2 instruction, loads data from shared memory
+ * to fragment
+ * \tparam T data type of the fragment
+ * \param R pointer to the fragment
+ * \param smem_ptr pointer to the shared memory
+ */
+ template <typename T>
+ __device__ __forceinline__ void ldmatrix_m8n8x2(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X2_ENABLED
+ uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+ asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
+ : "=r"(R[0]), "=r"(R[1])
+ : "r"(smem_int_ptr));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of PTX ldmatrix m8n8.x4 instruction, loads data from shared memory
+ * to fragment
+ * \tparam T data type of the fragment
+ * \param R pointer to the fragment
+ * \param smem_ptr pointer to the shared memory
+ */
+ template <typename T>
+ __device__ __forceinline__ void ldmatrix_m8n8x4(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X4_ENABLED
+ uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+ asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+ : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
+ : "r"(smem_int_ptr));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of PTX ldmatrix m8n8.x4 transposed instruction, loads data from
+ * shared memory to fragment and transposes the fragment
+ * \tparam T data type of the fragment
+ * \param R pointer to the fragment
+ * \param smem_ptr pointer to the shared memory
+ */
+ template <typename T>
+ __device__ __forceinline__ void ldmatrix_m8n8x4_trans(uint32_t* R, T* smem_ptr) {
+#ifdef LDMATRIX_M8N8X4_ENABLED
+ uint32_t smem_int_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
+ asm volatile("ldmatrix.sync.aligned.trans.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+ : "=r"(R[0]), "=r"(R[1]), "=r"(R[2]), "=r"(R[3])
+ : "r"(smem_int_ptr));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for ldmatrix instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
+ * multiplication, accumulated in f32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F32_M16N8K16_ENABLED
+ // ! only support half dtype now
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+ "f"(C[2]), "f"(C[3]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+ "f"(0.f), "f"(0.f));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
+ * multiplication, accumulated in f32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F32_M16N8K16_ENABLED
+ // ! only support half dtype now
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+ "f"(C[2]), "f"(C[3]));
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
+ "f"(C[6]), "f"(C[7]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+ "f"(0.f), "f"(0.f));
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
+ "f"(0.f), "f"(0.f));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n8k16 instruction for row major and column major f16 matrix
+ * multiplication, accumulated in f16.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n8k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F16_M16N8K16_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C[0]), "=r"(C[1])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C[0]), "=r"(C[1])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n16k16 instruction for row major and column major f16 matrix
+ * multiplication, accumulated in f16.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n16k16_row_col_f16f16f16(uint32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F16F16F16_M16N8K16_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C[0]), "=r"(C[1])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]));
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[2]), "r"(C[3]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C[0]), "=r"(C[1])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
+ asm volatile(
+ "mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int8 matrix
+ * multiplication, accumulated in int32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S8S8S32_M16N8K32_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+ "r"(C[2]), "r"(C[3]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+ "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n16k32 instruction for row major and column major int8 matrix
+ * multiplication, accumulated in int32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_s8s8s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S8S8S32_M16N8K32_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+ "r"(C[2]), "r"(C[3]));
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
+ "r"(C[6]), "r"(C[7]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+ "r"(0), "r"(0));
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
+ "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n8k32 instruction for row major and column major int4 matrix
+ * multiplication, accumulated in int32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n8k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S4S4S32_M16N8K64_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+ "r"(C[2]), "r"(C[3]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+ "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n16k64 instruction for row major and column major int4 matrix
+ * multiplication, accumulated in int32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n16k64_row_col_s4s4s32(int32_t* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_S4S4S32_M16N8K64_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C[0]), "r"(C[1]),
+ "r"(C[2]), "r"(C[3]));
+ asm volatile(
+ "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C[4]), "r"(C[5]),
+ "r"(C[6]), "r"(C[7]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[0]), "=r"(C[1]), "=r"(C[2]), "=r"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0),
+ "r"(0), "r"(0));
+ asm volatile(
+ "mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=r"(C[4]), "=r"(C[5]), "=r"(C[6]), "=r"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0),
+ "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n8k32 instruction for row major and column major fp8 e4m3 matrix
+ * multiplication, accumulated in fp32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n8k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F8F8F32_M16N8K16_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+ "f"(C[2]), "f"(C[3]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+ "f"(0.f), "f"(0.f));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
+ * multiplication, accumulated in fp16.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f16(uint32_t* C_uint32, uint32_t* A, uint32_t* B) {
+ // uint32_t* C_uint32 = reinterpret_cast<uint32_t*>(C);
+#ifdef MMA_F8F8F16_M16N8K16_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C_uint32[0]), "=r"(C_uint32[1])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(C_uint32[0]), "r"(C_uint32[1]));
+
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C_uint32[2]), "=r"(C_uint32[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(C_uint32[2]), "r"(C_uint32[3]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C_uint32[0]), "=r"(C_uint32[1])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "r"(0), "r"(0));
+
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
+ "{%0, %1},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, %9};\n"
+ : "=r"(C_uint32[2]), "=r"(C_uint32[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "r"(0), "r"(0));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Wrapper of the mma m16n16k32 instruction for row major and column major fp8 matrix
+ * multiplication, accumulated in fp32.
+ * \tparam mma_mode The mode of mma instruction, either kInit or kInplaceUpdate
+ * \param C pointer to the accumulator
+ * \param A pointer to the fragment of matrix A
+ * \param B pointer to the fragment of matrix B
+ */
+ template <MMAMode mma_mode = MMAMode::kInplaceUpdate>
+ __device__ __forceinline__ void mma_sync_m16n16k32_row_col_f8f8f32(float* C, uint32_t* A, uint32_t* B) {
+#ifdef MMA_F8F8F32_M16N8K16_ENABLED
+ if constexpr (mma_mode == MMAMode::kInplaceUpdate) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(C[0]), "f"(C[1]),
+ "f"(C[2]), "f"(C[3]));
+
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(C[4]), "f"(C[5]),
+ "f"(C[6]), "f"(C[7]));
+ } else if constexpr (mma_mode == MMAMode::kInit) {
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[0]), "=f"(C[1]), "=f"(C[2]), "=f"(C[3])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[0]), "r"(B[1]), "f"(0.f), "f"(0.f),
+ "f"(0.f), "f"(0.f));
+
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, %1, %2, %3},"
+ "{%4, %5, %6, %7},"
+ "{%8, %9},"
+ "{%10, %11, %12, %13};\n"
+ : "=f"(C[4]), "=f"(C[5]), "=f"(C[6]), "=f"(C[7])
+ : "r"(A[0]), "r"(A[1]), "r"(A[2]), "r"(A[3]), "r"(B[2]), "r"(B[3]), "f"(0.f), "f"(0.f),
+ "f"(0.f), "f"(0.f));
+ }
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Use mma instructions to compute rowsum.
+ */
+ __device__ __forceinline__ void rowsum_f16f16f32(float* d, uint32_t* s) {
+#ifdef MMA_F16F16F32_M16N8K16_ENABLED
+ asm volatile(
+ "{\n"
+ "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "
+ "{%0, _, %1, _},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, 0., %9, 0.};\n"
+ "}\n"
+ : "=f"(d[0]), "=f"(d[1])
+ : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(1006648320), // 1006648320 packs two 1.0f in half precision
+ "r"(1006648320), "f"(d[0]), "f"(d[1]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ /*!
+ * \brief Use mma instructions to compute rowsum.
+ */
+ __device__ __forceinline__ void rowsum_f8f8f32(float* d, uint32_t* s) {
+#ifdef MMA_F8F8F32_M16N8K16_ENABLED
+ asm volatile(
+ "mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
+ "{%0, _, %1, _},"
+ "{%2, %3, %4, %5},"
+ "{%6, %7},"
+ "{%8, 0., %9, 0.};\n"
+ : "=f"(d[0]), "=f"(d[1])
+ : "r"(s[0]), "r"(s[1]), "r"(s[2]), "r"(s[3]), "r"(943208504), "r"(943208504), // 943208504 packs four 1.0f in e4m3
+ "f"(d[0]), "f"(d[1]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for mma instruction");
+#endif
+ }
+
+ } // namespace mma
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/numeric_conversion.cuh b/src/core/cuda/sageattention/numeric_conversion.cuh
new file mode 100644
index 0000000..066f7fe
--- /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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_fp8.h>
+#include <cuda_runtime.h>
+
+namespace sageattention {
+
+#if (__CUDACC_VER_MAJOR__ * 10000 + __CUDACC_VER_MINOR__ * 100 >= 120400)
+#if (!defined(__CUDA_ARCH__) || (__CUDA_ARCH__ >= 890))
+#define FP8_CAST_ENABLED
+#endif
+#endif
+
+#if defined(__CUDA_ARCH__)
+#define RUNTIME_ASSERT(x) __brkpt()
+#else
+
+#define RUNTIME_ASSERT(x) assert(0 && x)
+#endif
+
+ __device__ __forceinline__ void unpack_half2_from_uint32_to_float(float* dest, uint32_t source) {
+ uint16_t h0 = source & 0xFFFF;
+ uint16_t h1 = (source >> 16) & 0xFFFF;
+ asm("cvt.f32.f16 %0, %1;" : "=f"(dest[0]) : "h"(h0));
+ asm("cvt.f32.f16 %0, %1;" : "=f"(dest[1]) : "h"(h1));
+ }
+
+ __device__ __forceinline__ void floatx4_to_e4m3x4(uint32_t* dest, float* source0, float* source1) {
+#ifdef FP8_CAST_ENABLED
+ asm volatile(
+ "{\n"
+ ".reg .b16 lo;\n"
+ ".reg .b16 hi;\n"
+ "cvt.rn.satfinite.e4m3x2.f32 lo, %2, %1;\n"
+ "cvt.rn.satfinite.e4m3x2.f32 hi, %4, %3;\n"
+ "mov.b32 %0, {lo, hi};\n"
+ "}"
+ : "=r"(dest[0]) : "f"(source0[0]), "f"(source0[1]), "f"(source1[0]), "f"(source1[1]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+ }
+
+ __device__ __forceinline__ void floatx4_to_e5m2x4(uint32_t* dest, float* source0, float* source1) {
+#ifdef FP8_CAST_ENABLED
+ asm volatile(
+ "{\n"
+ ".reg .b16 lo;\n"
+ ".reg .b16 hi;\n"
+ "cvt.rn.satfinite.e5m2x2.f32 lo, %2, %1;\n"
+ "cvt.rn.satfinite.e5m2x2.f32 hi, %4, %3;\n"
+ "mov.b32 %0, {lo, hi};\n"
+ "}"
+ : "=r"(dest[0]) : "f"(source0[0]), "f"(source1[1]), "f"(source1[0]), "f"(source1[1]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+ }
+
+ __device__ __forceinline__ void halfx4_to_e4m3x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
+#ifdef FP8_CAST_ENABLED
+ asm volatile(
+ "{\n"
+ ".reg .b16 lo;\n"
+ ".reg .b16 hi;\n"
+ "cvt.rn.satfinite.e4m3x2.f16x2 lo, %1;\n"
+ "cvt.rn.satfinite.e4m3x2.f16x2 hi, %2;\n"
+ "mov.b32 %0, {lo, hi};\n"
+ "}"
+ : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+ }
+
+ __device__ __forceinline__ void halfx4_to_e5m2x4(uint32_t* dest, uint32_t* source0, uint32_t* source1) {
+#ifdef FP8_CAST_ENABLED
+ asm volatile(
+ "{\n"
+ ".reg .b16 lo;\n"
+ ".reg .b16 hi;\n"
+ "cvt.rn.satfinite.e5m2x2.f16x2 lo, %1;\n"
+ "cvt.rn.satfinite.e5m2x2.f16x2 hi, %2;\n"
+ "mov.b32 %0, {lo, hi};\n"
+ "}"
+ : "=r"(dest[0]) : "r"(source0[0]), "r"(source1[0]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+ }
+
+ __device__ __forceinline__ void e4m3x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
+#ifdef FP8_CAST_ENABLED
+ asm volatile(
+ "{\n"
+ ".reg .b16 lo, hi;\n"
+ "mov.b32 {lo, hi}, %2;\n"
+ "cvt.rn.f16x2.e4m3x2 %0, lo;\n"
+ "cvt.rn.f16x2.e4m3x2 %1, hi;\n"
+ "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+ }
+
+ __device__ __forceinline__ void e5m2x4_to_halfx4(uint32_t* dest0, uint32_t* dest1, uint32_t* source) {
+#ifdef FP8_CAST_ENABLED
+ asm volatile(
+ "{\n"
+ ".reg .b16 lo, hi;\n"
+ "mov.b32 {lo, hi}, %2;\n"
+ "cvt.rn.f16x2.e5m2x2 %0, lo;\n"
+ "cvt.rn.f16x2.e5m2x2 %1, hi;\n"
+ "}\n" : "=r"(dest0[0]), "=r"(dest1[0]) : "r"(source[0]));
+#else
+ RUNTIME_ASSERT("Unsupported CUDA architecture for FP8 CAST instruction");
+#endif
+ }
+
+ __device__ __forceinline__ int8_t float_to_int8_rn(float x) {
+ uint32_t dst;
+ asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
+ return reinterpret_cast<const int8_t&>(dst);
+ }
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sageattention/permuted_smem.cuh b/src/core/cuda/sageattention/permuted_smem.cuh
new file mode 100644
index 0000000..37ad770
--- /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 <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include "cp_async.cuh"
+#include "mma.cuh"
+
+namespace sageattention {
+
+ enum class SwizzleMode {
+ k32B, // for k32B mode, a line of shared memory must have 32B (16 half value)
+ k64B, // for k64B mode, a line of shared memory must have 64B (32 half value)
+ k128B, // 128B already spans all banks in shared memory. a line of shared memory can have multiple 128B.
+ };
+
+ // Use 128bit as the granularity to fetch/store data per thread to maximize memory bandwidth
+ using b128_t = uint4;
+
+ /*!
+ * \brief A stateless shared memory wrapper that uses templates to avoid runtime conditionals. It makes sure
+ * that access to consecutive rows idx in the same column idx will make full use of the shared memory bank through
+ * permutation in the granularity of 128bit.
+ *
+ * This struct treats all offsets to be the number of `b128_t` elements. It is designed to be stateless,
+ * meaning it does not maintain any information about the current pointer position. The offset returnd by
+ * the struct can be used to access the shared memory through the provided interface.
+ *
+ * The struct guarantees that the read to permuted offset (i, j) will be the value stored in permuted offset (i, j).
+ * We assume that shared memory operation operates on at least two consecutive 128-bit values in a row within a warp.
+ * Under this assumption, we do not permute for k32B mode.
+ */
+ template <SwizzleMode swizzle_mode, uint32_t stride>
+ struct smem_t {
+ // The base pointer.
+ b128_t* base;
+ // How many b128_t value a row contains
+ // uint32_t stride;
+
+ __device__ __forceinline__ smem_t() : base(nullptr) {}
+ template <typename T>
+ __device__ __forceinline__ smem_t(T* base) : base((b128_t*)base) {
+ if constexpr (swizzle_mode == SwizzleMode::k128B) {
+ static_assert(stride % 8 == 0, "Stride must be multiple of 8 for 128B swizzle mode");
+ } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+ static_assert(stride == 4, "Stride must be 4 for 64B swizzle mode");
+ } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+ static_assert(stride == 2, "Stride must be 2 for 32B swizzle mode");
+ } else {
+ static_assert(swizzle_mode != swizzle_mode, "Unsupported swizzle mode");
+ }
+ }
+
+ /*!
+ * \brief Set the base pointer.
+ */
+ template <typename T>
+ __device__ __forceinline__ void set_base(T* new_base) {
+ base = (b128_t*)new_base;
+ }
+
+ /*!
+ * \brief Compute the element offset given coordinates in a permuted shared memory.
+ * \param i The row index.
+ * \param j The column index.
+ */
+ static __device__ __forceinline__ uint32_t get_permuted_offset(const uint32_t& i, const uint32_t& j) {
+ if constexpr (swizzle_mode == SwizzleMode::k128B) {
+ return i * stride + (j ^ (i % 8));
+ } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+ return i * stride + (j ^ ((i / 2) % 4));
+ } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+ return i * stride + j;
+ }
+ }
+
+ /*!
+ * \tparam step_size The step size to advance the offset in the permuted shared memory.
+ * \param offset The current offset.
+ */
+ template <uint32_t step_size>
+ static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset) {
+ if constexpr (swizzle_mode == SwizzleMode::k128B) {
+ static_assert(step_size % 8 == 0,
+ "Unsupported step size");
+ return offset + step_size;
+ } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+ static_assert(step_size == 4, "Unsupported step size");
+ return offset + step_size;
+ } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+ static_assert(step_size == 2, "Unsupported step size");
+ return offset + step_size;
+ }
+ }
+
+ // ! use with care
+ template <uint32_t step_size>
+ static __device__ __forceinline__ uint32_t advance_offset_by_column(const uint32_t& offset, const uint32_t& step_idx) {
+ if constexpr (swizzle_mode == SwizzleMode::k128B) {
+ static_assert(step_size == 2 || step_size == 4 || step_size % 8 == 0,
+ "Unsupported step size");
+ if constexpr (step_size == 2) {
+ return (offset ^ (0x2 + (0x4 * (step_idx % 2 == 1)))) + (step_idx % 4 == 3) * 8;
+ } else if constexpr (step_size == 4) {
+ return (offset ^ 0x4) + (step_idx % 2 == 1) * 8;
+ } else {
+ // step_size % 8 == 0
+ return offset + step_size;
+ }
+ } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+ static_assert(step_size == 2 || step_size == 4, "Unsupported step size");
+ if constexpr (step_size == 2) {
+ return (offset ^ 0x2) + (step_idx % 2 == 1) * 4;
+ } else {
+ return offset + step_size;
+ }
+ } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+ return offset + step_size;
+ }
+ }
+
+ template <uint32_t step_size>
+ static __device__ __forceinline__ uint32_t advance_offset_by_row(const uint32_t& offset) {
+ if constexpr (swizzle_mode == SwizzleMode::k128B) {
+ static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
+ if constexpr (step_size == 4) {
+ return (offset ^ 0x4) + step_size * stride;
+ } else {
+ // step_size % 8 == 0
+ return offset + step_size * stride;
+ }
+ } else if constexpr (swizzle_mode == SwizzleMode::k64B) {
+ static_assert(step_size == 4 || step_size % 8 == 0, "Unsupported step size");
+ if constexpr (step_size == 4) {
+ return (offset ^ 0x2) + step_size * stride;
+ } else {
+ // step_size % 8 == 0
+ return offset + step_size * stride;
+ }
+ } else if constexpr (swizzle_mode == SwizzleMode::k32B) {
+ return offset + step_size * stride;
+ }
+ }
+
+ __device__ __forceinline__ void ldmatrix_m8n8x2(const uint32_t& offset, uint32_t* R) const {
+ b128_t* smem_ptr = base + offset;
+ mma::ldmatrix_m8n8x2(R, smem_ptr);
+ }
+
+ __device__ __forceinline__ void ldmatrix_m8n8x4(const uint32_t& offset, uint32_t* R) const {
+ b128_t* smem_ptr = base + offset;
+ mma::ldmatrix_m8n8x4(R, smem_ptr);
+ }
+
+ __device__ __forceinline__ void ldmatrix_m8n8x4_trans(const uint32_t& offset, uint32_t* R) const {
+ b128_t* smem_ptr = base + offset;
+ mma::ldmatrix_m8n8x4_trans(R, smem_ptr);
+ }
+
+ template <cp_async::SharedMemFillMode fill_mode, typename T>
+ __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr, bool predicate) const {
+ b128_t* smem_ptr = base + offset;
+ cp_async::pred_load_128b<cp_async::PrefetchMode::kPrefetch, fill_mode>(
+ smem_ptr, reinterpret_cast<const b128_t*>(gptr), predicate);
+ }
+
+ template <typename T>
+ __device__ __forceinline__ void load_128b_async(const uint32_t& offset, const T* gptr) const {
+ b128_t* smem_ptr = base + offset;
+ cp_async::load_128b<cp_async::PrefetchMode::kPrefetch>(smem_ptr, reinterpret_cast<const b128_t*>(gptr));
+ }
+
+ template <typename T>
+ __device__ __forceinline__ void store_128b(const uint32_t& offset, T* gptr) const {
+ *reinterpret_cast<b128_t*>(gptr) = *(base + offset);
+ }
+ };
+
+} // namespace sageattention
+
+#endif
diff --git a/src/core/cuda/sol_attention.cu b/src/core/cuda/sol_attention.cu
new file mode 100644
index 0000000..0b580ab
--- /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 <cuda_bf16.h>
+#include <algorithm>
+#include <climits>
+
+namespace sd_cuda {
+
+ static __device__ __forceinline__ float warp_reduce_sum(float x) {
+#pragma unroll
+ for (int offset = 16; offset > 0; offset >>= 1)
+ x += __shfl_xor_sync(0xffffffff, x, offset);
+ return x;
+ }
+
+ static constexpr int SOL_DIM = 128;
+ static constexpr int SOL_BLOCK = 64;
+
+ struct SolWorkspace {
+ size_t q, k, v, kc, vc, qm, stats, threshold, size;
+ int tokens, blocks;
+
+ explicit SolWorkspace(const ggml_tensor* dst) {
+ const auto* input = dst->src[0];
+ const size_t heads = input->ne[2] * input->ne[3];
+ tokens = GGML_PAD(input->ne[1], SOL_BLOCK);
+ blocks = GGML_PAD(tokens / SOL_BLOCK, SOL_BLOCK);
+ size = 0;
+ q = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+ k = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+ v = reserve(heads * tokens * SOL_DIM * sizeof(nv_bfloat16));
+ kc = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+ vc = reserve(heads * blocks * SOL_DIM * sizeof(nv_bfloat16));
+ qm = reserve(heads * (tokens / SOL_BLOCK) * SOL_DIM * sizeof(float));
+ stats = reserve(heads * 2 * SOL_DIM * sizeof(float));
+ threshold = reserve(heads * (tokens / SOL_BLOCK) * sizeof(float));
+ }
+
+ size_t reserve(size_t bytes) {
+ const size_t offset = size;
+ size += GGML_PAD(bytes, 256);
+ return offset;
+ }
+ };
+
+ static __global__ void sol_pack(const float* q, const float* k, const float* v, nv_bfloat16* qb, nv_bfloat16* kb, nv_bfloat16* vb, int length, int padded) {
+ const int index = blockIdx.x * blockDim.x + threadIdx.x;
+ if (index >= padded * SOL_DIM) {
+ return;
+ }
+ const int token = index / SOL_DIM;
+ const int channel = index % SOL_DIM;
+ const size_t src = size_t(blockIdx.y) * length * SOL_DIM + index;
+ const size_t dst = size_t(blockIdx.y) * padded * SOL_DIM + index;
+ qb[dst] = __float2bfloat16(token < length ? q[src] : 0.f);
+ kb[dst] = __float2bfloat16(token < length ? k[src] : 0.f);
+ vb[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token] =
+ __float2bfloat16(token < length ? v[src] : 0.f);
+ }
+
+ static __global__ void sol_summaries(const nv_bfloat16* q, const nv_bfloat16* k, const nv_bfloat16* v, nv_bfloat16* kc, nv_bfloat16* vc, float* qm, int length, int padded, int padded_blocks) {
+ const int channel = threadIdx.x;
+ const int block = blockIdx.x;
+ const int count = min(SOL_BLOCK, max(0, length - block * SOL_BLOCK));
+ float sum_q = 0.f, sum_k = 0.f, sum_v = 0.f;
+ for (int i = 0; i < count; ++i) {
+ const int token = block * SOL_BLOCK + i;
+ const size_t index = (size_t(blockIdx.y) * padded + token) * SOL_DIM + channel;
+ sum_q += __bfloat162float(q[index]);
+ sum_k += __bfloat162float(k[index]);
+ sum_v += __bfloat162float(v[(size_t(blockIdx.y) * SOL_DIM + channel) * padded + token]);
+ }
+ kc[(size_t(blockIdx.y) * padded_blocks + block) * SOL_DIM + channel] =
+ __float2bfloat16(count ? sum_k / count : 0.f);
+ vc[(size_t(blockIdx.y) * SOL_DIM + channel) * padded_blocks + block] = __float2bfloat16(sum_v);
+ if (count) {
+ qm[(size_t(blockIdx.y) * (padded / SOL_BLOCK) + block) * SOL_DIM + channel] = sum_q / count;
+ }
+ }
+
+ static __global__ void sol_statistics(const nv_bfloat16* kc, float* stats, int blocks, int padded_blocks) {
+ const int channel = threadIdx.x;
+ float sum = 0.f, sum_sq = 0.f;
+ for (int i = 0; i < blocks; ++i) {
+ const float value = __bfloat162float(kc[(size_t(blockIdx.x) * padded_blocks + i) * SOL_DIM + channel]);
+ sum += value;
+ sum_sq += value * value;
+ }
+ const float mean = sum / blocks;
+ stats[size_t(blockIdx.x) * 2 * SOL_DIM + channel] = mean;
+ stats[size_t(blockIdx.x) * 2 * SOL_DIM + SOL_DIM + channel] = fmaxf(sum_sq / blocks - mean * mean, 0.f);
+ }
+
+ static __global__ void sol_thresholds(const float* qm, const float* stats, float* threshold, int blocks, float scale, float tau) {
+ const int channel = threadIdx.x;
+ const float q = qm[(size_t(blockIdx.y) * blocks + blockIdx.x) * SOL_DIM + channel];
+ float mean = q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + channel];
+ float var = q * q * stats[size_t(blockIdx.y) * 2 * SOL_DIM + SOL_DIM + channel];
+ mean = warp_reduce_sum(mean);
+ var = warp_reduce_sum(var);
+ __shared__ float means[4], vars[4];
+ if (channel % 32 == 0) {
+ means[channel / 32] = mean;
+ vars[channel / 32] = var;
+ }
+ __syncthreads();
+ if (channel == 0) {
+ mean = means[0] + means[1] + means[2] + means[3];
+ var = vars[0] + vars[1] + vars[2] + vars[3];
+ threshold[size_t(blockIdx.y) * blocks + blockIdx.x] =
+ mean * scale + tau * sqrtf(fmaxf(var * scale * scale, 0.f) + 1.e-6f);
+ }
+ }
+
+ struct SolQTile {
+ uint32_t x[4];
+ };
+ struct SolKTile {
+ uint32_t x[2];
+ };
+ struct SolTile {
+ float x[4] = {};
+ static __device__ __forceinline__ int get_i(int i) {
+ return (i / 2) * 8 + threadIdx.x / 4;
+ }
+ static __device__ __forceinline__ int get_j(int i) {
+ return (threadIdx.x % 4) * 2 + i % 2;
+ }
+ };
+
+ static __device__ __forceinline__ void sol_load(SolQTile& tile, const nv_bfloat162* input, int stride) {
+ const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+ for (int i = 0; i < 4; ++i)
+ tile.x[i] = data[((i % 2) * 8 + threadIdx.x / 4) * stride + (i / 2) * 4 + threadIdx.x % 4];
+ }
+ static __device__ __forceinline__ void sol_load(SolKTile& tile, const nv_bfloat162* input, int stride) {
+ const auto* data = reinterpret_cast<const uint32_t*>(input);
+#pragma unroll
+ for (int i = 0; i < 2; ++i)
+ tile.x[i] = data[(threadIdx.x / 4) * stride + i * 4 + threadIdx.x % 4];
+ }
+ static __device__ __forceinline__ void sol_mma(SolTile& output, const SolQTile& q, const SolKTile& k) {
+#if __CUDA_ARCH__ >= 800
+ auto* dst = reinterpret_cast<uint32_t*>(output.x);
+ asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};"
+ : "+r"(dst[0]), "+r"(dst[1]), "+r"(dst[2]), "+r"(dst[3])
+ : "r"(q.x[0]), "r"(q.x[1]), "r"(q.x[2]), "r"(q.x[3]), "r"(k.x[0]), "r"(k.x[1]));
+#else
+ asm("trap;");
+#endif
+ }
+
+ static __device__ __forceinline__ void sol_scores(const SolQTile (&q)[8], const nv_bfloat16* k, SolTile (&scores)[8], float scale) {
+#pragma unroll
+ for (int n = 0; n < 8; ++n) {
+ scores[n] = SolTile{};
+#pragma unroll
+ for (int d = 0; d < 8; ++d) {
+ SolKTile kt;
+ sol_load(kt, reinterpret_cast<const nv_bfloat162*>(k + n * 8 * SOL_DIM + d * 16), SOL_DIM / 2);
+ sol_mma(scores[n], q[d], kt);
+ }
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ scores[n].x[i] *= scale;
+ }
+ }
+ }
+
+ template <bool PROXY>
+ static __device__ __forceinline__ void sol_update(SolTile (&scores)[8], SolTile (&output)[16], float (&row_max)[2], float (&row_sum)[2], nv_bfloat16* probabilities, const nv_bfloat16* values, int stride, int start, int length, const int* exact) {
+ float maximum[2] = {-INFINITY, -INFINITY};
+#pragma unroll
+ for (int n = 0; n < 8; ++n) {
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ const int column = n * 8 + SolTile::get_j(i);
+ const bool valid = PROXY ? ((start + column) * SOL_BLOCK < length && !exact[column])
+ : (start + column < length);
+ scores[n].x[i] = valid ? scores[n].x[i] : -INFINITY;
+ maximum[i / 2] = fmaxf(maximum[i / 2], scores[n].x[i]);
+ }
+ }
+ float alpha[2], sums[2] = {0.f, 0.f};
+#pragma unroll
+ for (int row = 0; row < 2; ++row) {
+ maximum[row] = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 1));
+ maximum[row] = fmaxf(maximum[row], __shfl_xor_sync(0xffffffff, maximum[row], 2));
+ const float next = fmaxf(row_max[row], maximum[row]);
+ alpha[row] = row_max[row] == next ? 1.f : exp2f(row_max[row] - next);
+ row_max[row] = next;
+ }
+#pragma unroll
+ for (int n = 0; n < 8; ++n) {
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ const int column = n * 8 + SolTile::get_j(i);
+ const float p = scores[n].x[i] == -INFINITY ? 0.f : exp2f(scores[n].x[i] - row_max[i / 2]);
+ const int count = PROXY ? min(SOL_BLOCK, max(0, length - (start + column) * SOL_BLOCK)) : 1;
+ sums[i / 2] += p * count;
+ probabilities[SolTile::get_i(i) * SOL_BLOCK + column] = __float2bfloat16(p);
+ }
+ }
+#pragma unroll
+ for (int row = 0; row < 2; ++row) {
+ sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 1);
+ sums[row] += __shfl_xor_sync(0xffffffff, sums[row], 2);
+ row_sum[row] = row_sum[row] * alpha[row] + sums[row];
+ }
+ __syncwarp();
+#pragma unroll
+ for (int n = 0; n < 16; ++n) {
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ output[n].x[i] *= alpha[i / 2];
+ }
+#pragma unroll
+ for (int d = 0; d < 4; ++d) {
+ SolQTile pt;
+ SolKTile vt;
+ sol_load(pt, reinterpret_cast<const nv_bfloat162*>(probabilities + d * 16), SOL_BLOCK / 2);
+ sol_load(vt, reinterpret_cast<const nv_bfloat162*>(values + n * 8 * stride + d * 16), stride / 2);
+ sol_mma(output[n], pt, vt);
+ }
+ }
+ __syncwarp();
+ }
+
+ static __global__ __launch_bounds__(128) void sol_forward(
+ const nv_bfloat16* q,
+ const nv_bfloat16* k,
+ const nv_bfloat16* v,
+ const nv_bfloat16* kc,
+ const nv_bfloat16* vc,
+ const float* thresholds,
+ float* dst,
+ int length,
+ int padded,
+ int padded_blocks,
+ int heads,
+ float scale) {
+#if __CUDA_ARCH__ >= 800
+ const int warp = threadIdx.y;
+ const int lane = threadIdx.x;
+ const int tid = warp * 32 + lane;
+ const int blocks = padded / SOL_BLOCK;
+ const int q_start = blockIdx.x * SOL_BLOCK;
+ const int q_len = min(SOL_BLOCK, length - q_start);
+ const size_t base = size_t(blockIdx.y) * padded * SOL_DIM;
+ SolQTile qt[8];
+#pragma unroll
+ for (int d = 0; d < 8; ++d) {
+ sol_load(qt[d], reinterpret_cast<const nv_bfloat162*>(q + base + (q_start + warp * 16) * SOL_DIM + d * 16), SOL_DIM / 2);
+ }
+ SolTile output[16];
+ float row_max[2] = {-INFINITY, -INFINITY};
+ float row_sum[2] = {0.f, 0.f};
+ __shared__ nv_bfloat16 probabilities[SOL_BLOCK * SOL_BLOCK];
+ __shared__ float routes[4][SOL_BLOCK];
+ __shared__ int exact[SOL_BLOCK];
+ const float threshold = thresholds[size_t(blockIdx.y) * blocks + blockIdx.x];
+ for (int group = 0; group < blocks; group += SOL_BLOCK) {
+ SolTile scores[8];
+ sol_scores(qt, kc + (size_t(blockIdx.y) * padded_blocks + group) * SOL_DIM, scores, scale);
+#pragma unroll
+ for (int n = 0; n < 8; ++n) {
+#pragma unroll
+ for (int column = 0; column < 2; ++column) {
+ float sum = scores[n].x[column] + scores[n].x[column + 2];
+ sum += __shfl_xor_sync(0xffffffff, sum, 4);
+ sum += __shfl_xor_sync(0xffffffff, sum, 8);
+ sum += __shfl_xor_sync(0xffffffff, sum, 16);
+ if (lane < 4) {
+ routes[warp][n * 8 + 2 * lane + column] = sum;
+ }
+ }
+ }
+ __syncthreads();
+ if (tid < SOL_BLOCK) {
+ const int block = group + tid;
+ const float score = (routes[0][tid] + routes[1][tid] + routes[2][tid] + routes[3][tid]) / q_len;
+ exact[tid] = block < blocks && (score > threshold || abs(int(blockIdx.x) - block) <= 1);
+ }
+ __syncthreads();
+ sol_update<true>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+ vc + size_t(blockIdx.y) * padded_blocks * SOL_DIM + group,
+ padded_blocks, group, length, exact);
+ for (int block = 0; block < SOL_BLOCK && group + block < blocks; ++block) {
+ if (!exact[block]) {
+ continue;
+ }
+ const int start = (group + block) * SOL_BLOCK;
+ sol_scores(qt, k + base + start * SOL_DIM, scores, scale);
+ sol_update<false>(scores, output, row_max, row_sum, probabilities + warp * 16 * SOL_BLOCK,
+ v + base + start, padded, start, length, nullptr);
+ }
+ __syncthreads();
+ }
+ const int head = blockIdx.y % heads;
+ const int batch = blockIdx.y / heads;
+#pragma unroll
+ for (int n = 0; n < 16; ++n) {
+#pragma unroll
+ for (int i = 0; i < 4; ++i) {
+ const int token = q_start + warp * 16 + SolTile::get_i(i);
+ if (token < length) {
+ const int channel = n * 8 + SolTile::get_j(i);
+ dst[((size_t(batch) * length + token) * heads + head) * SOL_DIM + channel] =
+ __bfloat162float(__float2bfloat16(output[n].x[i] / row_sum[i / 2]));
+ }
+ }
+ }
+#else
+ asm("trap;");
+#endif
+ }
+
+ static bool sol_supports(ggml_backend_dev_t, const ggml_tensor* op, void*) {
+ SDSolAttentionParams params;
+ if (!sd_ggml_sol_attn_describe(op, ¶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<uintptr_t>(tensor->data) & 3)
+ return false;
+ return !(reinterpret_cast<uintptr_t>(op->data) & 3);
+ }
+
+ static ggml_status sol_prepare(ggml_backend_t, ggml_tensor* const* nodes, int, void**, size_t* workspace, void*) {
+ *workspace = SolWorkspace(nodes[0]).size;
+ return GGML_STATUS_SUCCESS;
+ }
+ static void sol_destroy(void*, void*) {}
+
+ static ggml_status sol_execute(const ggml_ext_execution_context* execution, ggml_tensor* const* nodes, int, void*, void*) {
+ ggml_ext_gpu_context context;
+ if (!cuda_context(execution->backend, &context))
+ return GGML_STATUS_FAILED;
+ auto* dst = nodes[0];
+ SDSolAttentionParams params;
+ if (!sd_ggml_sol_attn_describe(dst, ¶ms))
+ return GGML_STATUS_FAILED;
+ const auto* q = dst->src[0];
+ const SolWorkspace work(dst);
+ char* data = static_cast<char*>(ggml_backend_buffer_get_base(execution->workspace));
+ auto* qb = reinterpret_cast<nv_bfloat16*>(data + work.q);
+ auto* kb = reinterpret_cast<nv_bfloat16*>(data + work.k);
+ auto* vb = reinterpret_cast<nv_bfloat16*>(data + work.v);
+ auto* kc = reinterpret_cast<nv_bfloat16*>(data + work.kc);
+ auto* vc = reinterpret_cast<nv_bfloat16*>(data + work.vc);
+ auto* qm = reinterpret_cast<float*>(data + work.qm);
+ auto* stats = reinterpret_cast<float*>(data + work.stats);
+ auto* threshold = reinterpret_cast<float*>(data + work.threshold);
+ const int length = q->ne[1], heads = q->ne[2], batch_heads = heads * q->ne[3];
+ const int blocks = work.tokens / SOL_BLOCK;
+ const float scale = params.scale * 1.4426950408889634f;
+ const auto stream = static_cast<cudaStream_t>(context.stream);
+ sol_pack<<<dim3((work.tokens * SOL_DIM + 255) / 256, batch_heads), 256, 0, stream>>>(
+ static_cast<const float*>(q->data), static_cast<const float*>(dst->src[1]->data), static_cast<const float*>(dst->src[2]->data),
+ qb, kb, vb, length, work.tokens);
+ sol_summaries<<<dim3(work.blocks, batch_heads), SOL_DIM, 0, stream>>>(qb, kb, vb, kc, vc, qm, length, work.tokens, work.blocks);
+ sol_statistics<<<batch_heads, SOL_DIM, 0, stream>>>(kc, stats, blocks, work.blocks);
+ sol_thresholds<<<dim3(blocks, batch_heads), SOL_DIM, 0, stream>>>(qm, stats, threshold, blocks, scale, params.tau);
+ sol_forward<<<dim3(blocks, batch_heads), dim3(32, 4), 0, stream>>>(qb, kb, vb, kc, vc, threshold,
+ static_cast<float*>(dst->data), length, work.tokens, work.blocks, heads, scale);
+ return cudaGetLastError() == cudaSuccess ? GGML_STATUS_SUCCESS : GGML_STATUS_FAILED;
+ }
+
+ void register_sol_attention(ggml_backend_dev_t device) {
+ const auto info = cuda_device_info(device);
+ if (info.device < 0)
+ return;
+ int previous;
+ cudaDeviceProp properties;
+ cudaFuncAttributes attributes;
+ if (cudaGetDeviceProperties(&properties, info.device) != cudaSuccess || properties.major < 8 ||
+ cudaGetDevice(&previous) != cudaSuccess || cudaSetDevice(info.device) != cudaSuccess)
+ return;
+ const auto error = cudaFuncGetAttributes(&attributes, sol_forward);
+ if (cudaSetDevice(previous) != cudaSuccess || error != cudaSuccess ||
+ std::min(attributes.binaryVersion, attributes.ptxVersion) < 80)
+ return;
+ ggml_ext_kernel_desc desc = {};
+ desc.abi_version = GGML_EXT_ABI_VERSION;
+ desc.name = "sd.sol_attn.cuda.v1";
+ desc.op = GGML_OP_CUSTOM;
+ desc.custom_op = sd_ggml_sol_attn_op();
+ desc.priority = 100;
+ desc.supports = sol_supports;
+ desc.prepare = sol_prepare;
+ desc.execute = sol_execute;
+ desc.destroy = sol_destroy;
+ GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
+ }
+
+}
+#else
+namespace sd_cuda {
+ void register_sol_attention(ggml_backend_dev_t) {}
+}
+#endif
diff --git a/src/core/ggml_compat.cpp b/src/core/ggml_compat.cpp
new file mode 100644
index 0000000..061b337
--- /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 <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+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<const ggml_tensor*, ggml_tensor*> mapping;
+ std::unordered_set<ggml_tensor*> mapped;
+ std::vector<const ggml_tensor*> 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<ggml_tensor*> 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<std::pair<ggml_tensor*, int32_t>> 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<ggml_tensor*> 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<std::pair<ggml_tensor*, int32_t>> 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 0000000..9f64262
--- /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 5747d62..4a97704 100644
--- a/src/core/ggml_extend.h
+++ b/src/core/ggml_extend.h
@@ -4,6 +4,7 @@
#include <cstdint>
#include <vector>
+#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 8586d87..47fe80f 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 <algorithm>
#include <cctype>
@@ -17,9 +18,8 @@
#include <hip/hip_runtime_api.h>
#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 0000000..1f0aa1d
--- /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 <algorithm>
+#include <cmath>
+#include <cstring>
+#include <mutex>
+#include <unordered_set>
+#include <utility>
+
+#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<ggml_backend_dev_t> registered;
+ std::lock_guard<std::mutex> 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 0000000..49b0e76
--- /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 8454d43..fc6d8c6 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<uint64_t> graph_layout(ggml_cgraph* graph, bool include_bindings) {
- const size_t graph_size = static_cast<size_t>(graph->n_leafs) + graph->n_nodes;
+ const size_t graph_size = static_cast<size_t>(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<uint64_t>& layout) {
GGML_ASSERT(gf != nullptr);
- if (plan.leaf_names.size() != static_cast<size_t>(gf->n_leafs) ||
+ if (plan.leaf_names.size() != static_cast<size_t>(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 0000000..8e55d99
--- /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 <cstdint>\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 0000000..4caadfb
--- /dev/null
+++ b/src/core/vulkan/int8.cpp
@@ -0,0 +1,806 @@
+#include <algorithm>
+#include <array>
+#include <climits>
+#include <cstring>
+#include <vector>
+#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_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1"));
+ auto compute = reinterpret_cast<ggml_ext_vulkan_compute_api_fn>(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<ggml_cgraph*>(graph), i);
+ }
+ static int count_nodes(const ggml_cgraph* graph) {
+ return ggml_graph_n_nodes(const_cast<ggml_cgraph*>(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<uintptr_t>(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 <size_t N>
+ 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<ggml_ext_vulkan_binding>& 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<ggml_ext_vulkan_binding> 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<ggml_ext_vulkan_binding> 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<ggml_ext_vulkan_binding> 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<uint32_t> 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<uintptr_t>(output->data);
+ const auto src = reinterpret_cast<uintptr_t>(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 0000000..63b976e
--- /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 <algorithm>
+#include <climits>
+#include <cstring>
+
+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_ext_vulkan_api_fn>(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<uintptr_t>(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<ggml_cgraph*>(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 0000000..3e4b3a2
--- /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 0000000..7c0292e
--- /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 0000000..573e257
--- /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<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator> sums[TILE_M * TILE_N];
+ [[unroll]]
+ for (uint i = 0; i < TILE_M * TILE_N; ++i) {
+ sums[i] = coopmat<int32_t, gl_ScopeSubgroup, CM_M, CM_N, gl_MatrixUseAccumulator>(0);
+ }
+
+ for (uint k_base = 0; k_base < p.k; k_base += CM_K) {
+ coopmat<int8_t, gl_ScopeSubgroup, CM_M, CM_K, gl_MatrixUseA> activations[TILE_M];
+ coopmat<int8_t, gl_ScopeSubgroup, CM_K, CM_N, gl_MatrixUseB> 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 0000000..3cb4547
--- /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<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(0);
+ coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator> sum_up = coopmat<int32_t, gl_ScopeWorkgroup, BM, BN, gl_MatrixUseAccumulator>(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<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+ coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> 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<int8_t, gl_ScopeWorkgroup, BM, BK, gl_MatrixUseA> w;
+ coopmat<int8_t, gl_ScopeWorkgroup, BK, BN, gl_MatrixUseB> 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 0000000..20edd61
--- /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 0000000..eb451f6
--- /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 0000000..7d08f74
--- /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 0000000..ccfc6d6
--- /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 63584bc..b6fdd32 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 739ef77..3dd720c 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 0000000..ad7b2f5
--- /dev/null
+++ b/thirdparty/cutlass
@@ -0,0 +1 @@
+Subproject commit ad7b2f5e84fcfa124cb02b91d5bd26d238c0459e