DarthAffe/stable-diffusion.cpp
1
0 Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-11-12 11:56:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: preserve INT8 APIs through backend extensions

Browse Source
This commit is contained in:
leejet committed 2026-11-01 01:55:51 +08:00
1 parent 051b49751e
commit f0204a7eaa
20 files changed
+1579 -108
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/core/cpu/int8.cpp
-8
View File
@@ -168,10 +168,6 @@ namespace sd_cpu {
d.prepare = i8_prepare;
d.execute = i8_execute;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.name = "sd.linear_i8.cpu.override.v1";
d.op = GGML_OP_MUL_MAT;
d.custom_op = nullptr;
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();
@@ -179,10 +175,6 @@ namespace sd_cpu {
d.prepare = convrot_prepare;
d.execute = convrot_execute;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.name = "sd.quantize_i8_convrot.cpu.override.v1";
d.op = GGML_OP_QUANTIZE_I8_CONVROT;
d.custom_op = nullptr;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
}
}
src/core/cuda/cute_attention.cu
+11 -7
View File
@@ -46,6 +46,7 @@
#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) {
@@ -549,6 +550,7 @@ static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, in
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) {
@@ -558,8 +560,8 @@ static int output_fusion(ggml_backend_dev_t device, const ggml_cgraph* graph, in
if (!ggml_is_contiguous(node) || !ggml_are_same_shape(previous, node))
break;
scaled = true;
} else if (node->op == GGML_OP_QUANTIZE_I8_CONVROT) {
if (node->op_params[0] == 256 && previous->ne[0] % 256 == 0 &&
} 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;
@@ -579,7 +581,7 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
if (!fusion_enabled())
return 0;
int offset = 0;
if (graph_node(graph, first)->op == GGML_OP_MUL_MAT) {
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)) ||
@@ -627,8 +629,8 @@ struct AttentionPlan {
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 = nodes[0]->op == GGML_OP_MUL_MAT ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
: 0;
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;
@@ -738,7 +740,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
}
}
auto* dst = nodes[count - 1];
if (dst->op == GGML_OP_QUANTIZE_I8_CONVROT) {
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)
@@ -782,7 +785,8 @@ void sd_ggml_register_cuda_attention(ggml_backend_dev_t device) {
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_MUL_MAT;
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));
}
src/core/cuda/i8_gated_residual.cu
+2 -6
View File
@@ -449,7 +449,8 @@ 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_MUL_MAT;
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;
@@ -457,9 +458,4 @@ void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
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));
desc.name = "sd.custom_i8_gated_residual.cuda.v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_i8_op();
desc.optimize = nullptr;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
src/core/cuda/i8_projection.cu
+10 -5
View File
@@ -1,5 +1,6 @@
#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>
@@ -188,10 +189,13 @@ static bool aligned(const ggml_tensor* tensor) {
}
static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) {
const auto* w = mat->src[0];
const auto* x = mat->src[1];
const auto* ws = mat->src[2];
if (mat->op != GGML_OP_MUL_MAT || mat->type != GGML_TYPE_F32 || mat->op_params[4] <= 0 || !w || !x || !ws || mat->src[3] ||
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) ||
@@ -379,7 +383,8 @@ 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_MUL_MAT;
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_i8_op();
desc.priority = 350;
desc.match = projection_match;
desc.prepare = projection_prepare;
src/core/cuda/int8.cu
+2 -14
View File
@@ -462,8 +462,8 @@ namespace SD_CUDA_NAMESPACE {
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;
// Custom GEMM previously used cuBLAS; retain that path when I32 accumulation can overflow.
if (op.weight_scale && !op.bias && (nodes[0]->op == GGML_OP_MUL_MAT || p->k <= INT_MAX / 16384) &&
// 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)) {
@@ -723,29 +723,17 @@ namespace SD_CUDA_NAMESPACE {
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." SD_CUDA_NAME ".override.v1";
desc.op = GGML_OP_MUL_MAT;
desc.custom_op = nullptr;
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.name = "sd.linear_i8.epilogue." SD_CUDA_NAME ".custom.v1";
desc.op = GGML_OP_CUSTOM;
desc.custom_op = sd_ggml_i8_op();
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));
desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".override.v1";
desc.op = GGML_OP_QUANTIZE_I8_CONVROT;
desc.custom_op = nullptr;
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
}
}
src/core/cuda/rope.cu
+1 -1
View File
@@ -257,7 +257,7 @@ namespace SD_CUDA_NAMESPACE {
if (j >= 2) {
const auto* mat = ggml_graph_node(graph, j - 2);
const auto* view = ggml_graph_node(graph, j - 1);
if (mat->op == GGML_OP_MUL_MAT && mat->src[0]->type == GGML_TYPE_I8 &&
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;
src/core/ggml_extensions.cpp
+13 -30
View File
@@ -161,14 +161,7 @@ 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 (node->op == GGML_OP_MUL_MAT) {
op->weight_scale = node->src[2];
op->bias = node->src[3];
op->group_size = node->op_params[2];
op->packed = node->op_params[4] > 0;
if (op->packed)
op->group_size = node->op_params[4];
} else if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
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)))
@@ -191,16 +184,23 @@ bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
if (!node || !node->src[0] || !ggml_is_contiguous(node))
return false;
if (node->op == GGML_OP_QUANTIZE_I8_CONVROT)
*group = node->op_params[0];
else if (ggml_ext_op_get(node) == sd_ggml_convrot_op())
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));
}
static ggml_tensor* i8_create(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* xs, ggml_tensor* ws, ggml_tensor* bias, int group, const ggml_tensor* logical) {
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;
@@ -218,24 +218,7 @@ static ggml_tensor* i8_create(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x,
}
return ggml_ext_op_create(ctx, sd_ggml_i8_op(), inputs, count, &attr);
}
ggml_tensor* sd_ggml_mul_mat_i8(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* xs, ggml_tensor* ws, ggml_tensor* bias) {
if (!w || !x || !xs || x->type != GGML_TYPE_I8)
return nullptr;
return i8_create(ctx, w, x, xs, ws, bias, 0, x);
}
ggml_tensor* sd_ggml_linear_i8(ggml_context* ctx, ggml_tensor* w, ggml_tensor* x, ggml_tensor* ws, ggml_tensor* bias, int group) {
if (!w || !x)
return nullptr;
const auto* logical = x;
if (x->type == GGML_TYPE_I8) {
int packed_group;
if (!sd_ggml_convrot_describe(x, &packed_group) || packed_group != group)
return nullptr;
logical = x->src[0];
}
return i8_create(ctx, w, x, nullptr, ws, bias, group, logical);
}
ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
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);
}
src/core/ggml_extensions.h
+2 -3
View File
@@ -32,9 +32,8 @@ 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* sd_ggml_mul_mat_i8(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* input_scale, ggml_tensor* weight_scale = nullptr, ggml_tensor* bias = nullptr);
ggml_tensor* sd_ggml_linear_i8(ggml_context* ctx, ggml_tensor* weight, ggml_tensor* input, ggml_tensor* weight_scale, ggml_tensor* bias, int group_size);
ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, 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);
#endif
#endif
src/core/vulkan/CMakeLists.txt
+61 -19
View File
@@ -1,27 +1,69 @@
find_package(Vulkan COMPONENTS glslc REQUIRED)
set(SD_VULKAN_GENERATED "${CMAKE_CURRENT_BINARY_DIR}/shaders")
file(MAKE_DIRECTORY "${SD_VULKAN_GENERATED}")
foreach(NAME rope_apply rms_rope_apply rms_rope_cast)
set(SHADER "${NAME}")
set(FLAGS)
if(NAME STREQUAL "rms_rope_cast")
set(SHADER rms_rope_apply)
set(FLAGS -DCAST_F16)
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()
set(SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${SHADER}.comp")
set(SPV "${SD_VULKAN_GENERATED}/${NAME}.spv")
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")
# Keep native RoPE rounding: spirv-opt changes the subgroup reduction results.
add_custom_command(OUTPUT "${HEADER}"
COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=vulkan1.2 ${FLAGS} "${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)
list(APPEND SD_VULKAN_HEADERS "${HEADER}")
endforeach()
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")
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}")
src/core/vulkan/int8.cpp
+771
View File
@@ -0,0 +1,771 @@
#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) {
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;
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))
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 ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) {
*bytes = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) {
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));
}
}
src/core/vulkan/shaders/modulate_convrot.comp
+149
View File
@@ -0,0 +1,149 @@
#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();
}
}
src/core/vulkan/shaders/mul_mat_i8_tensorwise.comp
+62
View File
@@ -0,0 +1,62 @@
#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;
}
}
src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm1.comp
+94
View File
@@ -0,0 +1,94 @@
#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);
}
}
}
src/core/vulkan/shaders/mul_mat_i8_tensorwise_cm2.comp
+153
View File
@@ -0,0 +1,153 @@
#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);
}
}
src/core/vulkan/shaders/quantize_i8_convrot.comp
+102
View File
@@ -0,0 +1,102 @@
#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);
}
}
src/core/vulkan/shaders/quantize_i8_convrot_subgroup.comp
+117
View File
@@ -0,0 +1,117 @@
#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();
}
}