mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 11:56:36 +00:00
refactor: preserve INT8 APIs through backend extensions
This commit is contained in:
1 parent
051b49751e
commit
f0204a7eaa
20 files changed
+1573
-102
No files matched your search
+1
-1
@@ -374,7 +374,7 @@ if (NOT SD_USE_UPSTREAM_GGML AND (SD_CUDA OR SD_HIPBLAS OR SD_MUSA))
|
|||||||
endif()
|
endif()
|
||||||
if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
|
if (SD_VULKAN AND NOT SD_USE_UPSTREAM_GGML)
|
||||||
add_subdirectory(src/core/vulkan)
|
add_subdirectory(src/core/vulkan)
|
||||||
list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope)
|
list(APPEND SD_BACKEND_REGISTRATIONS sd_ggml_register_vulkan_rope sd_ggml_register_vulkan_i8)
|
||||||
endif()
|
endif()
|
||||||
if(SD_BACKEND_REGISTRATIONS)
|
if(SD_BACKEND_REGISTRATIONS)
|
||||||
string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")
|
string(REPLACE ";" ",\n " SD_BACKEND_REGISTRATIONS "${SD_BACKEND_REGISTRATIONS}")
|
||||||
|
|||||||
@@ -4,6 +4,7 @@
|
|||||||
#include "core/cuda/registration.h"
|
#include "core/cuda/registration.h"
|
||||||
|
|
||||||
void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
|
void sd_ggml_register_vulkan_rope(ggml_backend_dev_t device);
|
||||||
|
void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device);
|
||||||
|
|
||||||
static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = {
|
static void (*const sd_backend_extensions[])(ggml_backend_dev_t) = {
|
||||||
@SD_BACKEND_REGISTRATIONS@
|
@SD_BACKEND_REGISTRATIONS@
|
||||||
|
|||||||
+26
-13
@@ -190,22 +190,27 @@ and matchers now live in sd.cpp. `GGML_CUDA_FORCE_CUBLAS` disables these epilogu
|
|||||||
|
|
||||||
`src/core/ggml_extensions.cpp` defines `sd.linear_i8.v1` and
|
`src/core/ggml_extensions.cpp` defines `sd.linear_i8.v1` and
|
||||||
`sd.quantize_i8_convrot.v1`. `src/core/cpu/int8.cpp` supplies their CPU
|
`sd.quantize_i8_convrot.v1`. `src/core/cpu/int8.cpp` supplies their CPU
|
||||||
implementations and overrides of the old builtins. The CPU kernels use the
|
implementations. The CPU kernels use the
|
||||||
configured thread pool and scheduler workspace; the old CPU INT8 GEMM and ConvRot
|
configured thread pool and scheduler workspace; the old CPU INT8 GEMM and ConvRot
|
||||||
implementations have been removed from ggml.
|
implementations have been removed from ggml.
|
||||||
Custom nodes execute on CPU or CUDA when supported. Model code retains builtin
|
Model code keeps calling `ggml_ext_linear_i8_tensorwise`,
|
||||||
nodes while the existing CUDA and Vulkan fusion rules are being migrated.
|
`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, `sd_ggml_mul_mat_i8` accepts ordinary I8 weight
|
For already quantized activations, `ggml_mul_mat_i8_tensorwise` accepts ordinary
|
||||||
and activation tensors, a separate F32 activation scale tensor with one value per
|
I8 weight and activation tensors. Its optional final `input_scale` parameter is a
|
||||||
row, optional F32 weight scales with one value per output channel, and optional
|
separate F32 tensor with one value per activation row. Weight scales have one
|
||||||
bias. It computes `float(int32_dot(W, X)) * X_scale * W_scale + bias`, in that
|
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
|
order. All scales are explicit graph dependencies; there is no quantized-tensor
|
||||||
wrapper. Shapes use ggml's low-to-high dimension order. Call
|
wrapper. Shapes use ggml's low-to-high dimension order. Call
|
||||||
`sd_ggml_register_extensions` before scheduling these custom nodes.
|
`sd_ggml_register_extensions` before scheduling these custom nodes.
|
||||||
|
|
||||||
`sd_ggml_linear_i8` also accepts F32 activations (quantized into temporary
|
`ggml_mul_mat_i8_tensorwise` also accepts F32 activations (quantized into temporary
|
||||||
workspace) or the output of `sd_ggml_quantize_i8_convrot`. ConvRot keeps the existing
|
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
|
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.
|
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
|
Unsupported layouts and malformed scale shapes are rejected. Separate CPU quantize
|
||||||
@@ -217,7 +222,6 @@ ConvRot, and rescaling kernels. F32 inputs are quantized into workspace; ordinar
|
|||||||
I8 inputs retain separate scale tensors and use a padded temporary only when
|
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
|
needed by cuBLAS. Packed ConvRot inputs are consumed directly. Large ConvRot rows
|
||||||
use multiple kernels when they exceed shared memory. Aliasing outputs use scratch.
|
use multiple kernels when they exceed shared memory. Aliasing outputs use scratch.
|
||||||
The registrations also intercept builtin INT8 GEMM and standalone ConvRot nodes.
|
|
||||||
CUDA GEMM requires Turing or newer, contiguous tensors, K/N divisible by four,
|
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;
|
float-aligned inputs, and row counts that fit a padded signed 32-bit dimension;
|
||||||
CUDA ConvRot uses groups of 256.
|
CUDA ConvRot uses groups of 256.
|
||||||
@@ -227,7 +231,7 @@ 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,
|
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
|
small-row policy, Stream-K scheduling, Ada column bands, and separate scaling
|
||||||
order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row
|
order. The same extension matches `GEMM -> SwiGLU [-> ConvRot]` and row
|
||||||
concatenation for builtin and custom GEMM nodes. Weight scales remain explicit
|
concatenation for custom GEMM nodes. Weight scales remain explicit
|
||||||
inputs, including the gate/up scale permutation recorded on every SwiGLU launch.
|
inputs, including the gate/up scale permutation recorded on every SwiGLU launch.
|
||||||
Plans reserve bounded Stream-K workspace, cache occupancy on first execution,
|
Plans reserve bounded Stream-K workspace, cache occupancy on first execution,
|
||||||
and initialize barriers on every launch, so CUDA Graph replay observes updated
|
and initialize barriers on every launch, so CUDA Graph replay observes updated
|
||||||
@@ -248,7 +252,7 @@ implementations remain CUDA-specific.
|
|||||||
|
|
||||||
`src/core/cuda/i8_gated_residual.cu` implements
|
`src/core/cuda/i8_gated_residual.cu` implements
|
||||||
`INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered
|
`INT8 GEMM -> [Scale] -> Mul(gate) -> Add(residual) -> [Concat]` as a registered
|
||||||
CUDA fusion, for both builtin GEMM and custom I8 nodes with separate activation
|
CUDA fusion, for packed ConvRot inputs and ordinary I8 nodes with separate activation
|
||||||
scales. The corresponding native CUDA fusion rules and CUTLASS residual epilogue
|
scales. The corresponding native CUDA fusion rules and CUTLASS residual epilogue
|
||||||
have been removed from ggml. Its optimizer moves independent gate producers ahead
|
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.
|
of GEMM and optional scaling through `ggml_ext_graph_delay`, before fusion matching.
|
||||||
@@ -267,7 +271,7 @@ disable settings are honored.
|
|||||||
`src/core/cuda/convrot.cu` registers CUDA fusion for `Scale -> ConvRot`,
|
`src/core/cuda/convrot.cu` registers CUDA fusion for `Scale -> ConvRot`,
|
||||||
`Norm -> Mul -> ConvRot`, `RMSNorm -> Mul(weight) -> Mul(modulation) -> Add ->
|
`Norm -> Mul -> ConvRot`, `RMSNorm -> Mul(weight) -> Mul(modulation) -> Add ->
|
||||||
Add(shift) -> ConvRot`, `SwiGLU -> ConvRot`, and `SiLU/Sigmoid -> Mul -> ConvRot`.
|
Add(shift) -> ConvRot`, `SwiGLU -> ConvRot`, and `SiLU/Sigmoid -> Mul -> ConvRot`.
|
||||||
Both builtin and custom ConvRot outputs are accepted. SwiGLU accepts packed,
|
ConvRot outputs are custom nodes. SwiGLU accepts packed,
|
||||||
swapped, or separate gate/up inputs. Independent modulation and up-branch
|
swapped, or separate gate/up inputs. Independent modulation and up-branch
|
||||||
producers are moved ahead of their consumers through the safe graph-delay API.
|
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
|
Reordering keeps an owning producer together with in-place consumers, such as
|
||||||
@@ -399,3 +403,12 @@ Remaining model-specific CUDA paths and Vulkan INT8/ConvRot implementations stil
|
|||||||
need auditing or migration.
|
need auditing or migration.
|
||||||
This implementation does not yet externalize all model-specific kernels or add
|
This implementation does not yet externalize all model-specific kernels or add
|
||||||
dynamic dtypes.
|
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.
|
||||||
+1
-1
Submodule ggml updated: 394d41dd40...9201d070ff.
@@ -168,10 +168,6 @@ namespace sd_cpu {
|
|||||||
d.prepare = i8_prepare;
|
d.prepare = i8_prepare;
|
||||||
d.execute = i8_execute;
|
d.execute = i8_execute;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
|
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.name = "sd.quantize_i8_convrot.cpu.v1";
|
||||||
d.op = GGML_OP_CUSTOM;
|
d.op = GGML_OP_CUSTOM;
|
||||||
d.custom_op = sd_ggml_convrot_op();
|
d.custom_op = sd_ggml_convrot_op();
|
||||||
@@ -179,10 +175,6 @@ namespace sd_cpu {
|
|||||||
d.prepare = convrot_prepare;
|
d.prepare = convrot_prepare;
|
||||||
d.execute = convrot_execute;
|
d.execute = convrot_execute;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
|
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));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -46,6 +46,7 @@
|
|||||||
#include <unordered_map>
|
#include <unordered_map>
|
||||||
#include "core/cuda/i8_projection.h"
|
#include "core/cuda/i8_projection.h"
|
||||||
#include "core/cuda/rope.h"
|
#include "core/cuda/rope.h"
|
||||||
|
#include "core/ggml_extensions.h"
|
||||||
#include "ggml-ext.h"
|
#include "ggml-ext.h"
|
||||||
|
|
||||||
static __device__ __forceinline__ float sd_warp_max(float v) {
|
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;
|
bool scaled = false;
|
||||||
for (int count = 1; count < 8 && first + count < total; ++count) {
|
for (int count = 1; count < 8 && first + count < total; ++count) {
|
||||||
const auto* node = graph_node(graph, first + count);
|
const auto* node = graph_node(graph, first + count);
|
||||||
|
int convrot_group;
|
||||||
if (node->src[0] != previous)
|
if (node->src[0] != previous)
|
||||||
break;
|
break;
|
||||||
if (node->op == GGML_OP_VIEW || node->op == GGML_OP_RESHAPE) {
|
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))
|
if (!ggml_is_contiguous(node) || !ggml_are_same_shape(previous, node))
|
||||||
break;
|
break;
|
||||||
scaled = true;
|
scaled = true;
|
||||||
} else if (node->op == GGML_OP_QUANTIZE_I8_CONVROT) {
|
} else if (sd_ggml_convrot_describe(node, &convrot_group)) {
|
||||||
if (node->op_params[0] == 256 && previous->ne[0] % 256 == 0 &&
|
if (convrot_group == 256 && previous->ne[0] % 256 == 0 &&
|
||||||
size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
|
size_t(previous->ne[0]) * sizeof(float) + 128 <= cuda_device(device)->properties.sharedMemPerBlockOptin)
|
||||||
return count + 1;
|
return count + 1;
|
||||||
break;
|
break;
|
||||||
@@ -579,7 +581,7 @@ static int attention_match(ggml_backend_dev_t device, const ggml_cgraph* graph,
|
|||||||
if (!fusion_enabled())
|
if (!fusion_enabled())
|
||||||
return 0;
|
return 0;
|
||||||
int offset = 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)) ||
|
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),
|
!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 + 2), graph_node(graph, first + 3), graph_node(graph, first + 4)) ||
|
||||||
@@ -627,7 +629,7 @@ struct AttentionPlan {
|
|||||||
static ggml_status attention_prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void** state, size_t* workspace_size, void*) {
|
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{};
|
auto* plan = new AttentionPlan{};
|
||||||
*state = plan;
|
*state = plan;
|
||||||
plan->index = nodes[0]->op == GGML_OP_MUL_MAT ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
|
plan->index = ggml_ext_op_get(nodes[0]) == sd_ggml_i8_op() ? 5 : nodes[0]->op == GGML_OP_RMS_NORM ? 3
|
||||||
: 0;
|
: 0;
|
||||||
GGML_ASSERT(plan->index < count);
|
GGML_ASSERT(plan->index < count);
|
||||||
const auto* attn = nodes[plan->index];
|
const auto* attn = nodes[plan->index];
|
||||||
@@ -738,7 +740,8 @@ static ggml_status attention_execute(const ggml_ext_execution_context* execution
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
auto* dst = nodes[count - 1];
|
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;
|
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);
|
auto* packed = static_cast<int8_t*>(dst->data);
|
||||||
if (rows != padded_rows && cudaMemsetAsync(packed + width * rows, 0, width * (padded_rows - rows), stream) != cudaSuccess)
|
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;
|
desc.optimize = nullptr;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
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.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;
|
desc.priority = 400;
|
||||||
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
GGML_ASSERT(ggml_ext_register_kernel_ex(device, &desc, GGML_EXT_KERNEL_CAPTURE_SAFE | GGML_EXT_KERNEL_REPLAY_REUSABLE));
|
||||||
}
|
}
|
||||||
@@ -449,7 +449,8 @@ void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
|
|||||||
ggml_ext_kernel_desc desc = {};
|
ggml_ext_kernel_desc desc = {};
|
||||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||||
desc.name = "sd.i8_gated_residual.cuda.v1";
|
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.priority = 325;
|
||||||
desc.optimize = sd_i8_gated::optimize;
|
desc.optimize = sd_i8_gated::optimize;
|
||||||
desc.match = sd_i8_gated::match;
|
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.execute = sd_i8_gated::execute;
|
||||||
desc.destroy = sd_i8_gated::destroy;
|
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));
|
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));
|
|
||||||
}
|
}
|
||||||
@@ -1,5 +1,6 @@
|
|||||||
#include "core/cuda/i8_projection.h"
|
#include "core/cuda/i8_projection.h"
|
||||||
#include "core/cuda/rope.h"
|
#include "core/cuda/rope.h"
|
||||||
|
#include "core/ggml_extensions.h"
|
||||||
|
|
||||||
#include <cuda_fp16.h>
|
#include <cuda_fp16.h>
|
||||||
#include <cutlass/gemm/device/gemm_universal_adapter.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) {
|
static bool projection_supports(ggml_backend_dev_t device, const ggml_tensor* mat, const ggml_tensor* view) {
|
||||||
const auto* w = mat->src[0];
|
SDI8Operation op;
|
||||||
const auto* x = mat->src[1];
|
if (!sd_ggml_i8_describe(mat, &op) || !op.packed || !op.weight_scale || op.bias)
|
||||||
const auto* ws = mat->src[2];
|
return false;
|
||||||
if (mat->op != GGML_OP_MUL_MAT || mat->type != GGML_TYPE_F32 || mat->op_params[4] <= 0 || !w || !x || !ws || mat->src[3] ||
|
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->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] ||
|
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) ||
|
!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 = {};
|
ggml_ext_kernel_desc desc = {};
|
||||||
desc.abi_version = GGML_EXT_ABI_VERSION;
|
desc.abi_version = GGML_EXT_ABI_VERSION;
|
||||||
desc.name = "sd.i8_projection.cuda.v1";
|
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.priority = 350;
|
||||||
desc.match = projection_match;
|
desc.match = projection_match;
|
||||||
desc.prepare = projection_prepare;
|
desc.prepare = projection_prepare;
|
||||||
|
|||||||
+2
-14
@@ -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);
|
p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0);
|
||||||
#ifdef SD_CUDA_I8_CUTLASS
|
#ifdef SD_CUDA_I8_CUTLASS
|
||||||
size_t gemm_bytes = 0;
|
size_t gemm_bytes = 0;
|
||||||
// Custom GEMM previously used cuBLAS; retain that path when I32 accumulation can overflow.
|
// Retain cuBLAS when I32 accumulation can overflow.
|
||||||
if (op.weight_scale && !op.bias && (nodes[0]->op == GGML_OP_MUL_MAT || p->k <= INT_MAX / 16384) &&
|
if (op.weight_scale && !op.bias && p->k <= INT_MAX / 16384 &&
|
||||||
aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
|
aligned(op.weight, 16) && aligned(op.weight_scale, 16) && aligned(dst, 16) &&
|
||||||
(op.input->type == GGML_TYPE_F32 || (aligned(op.input, 16) && aligned(op.input_scale, 16))) &&
|
(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)) {
|
(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.execute = sd_i8_cuda::execute;
|
||||||
desc.destroy = sd_i8_cuda::destroy;
|
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));
|
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.name = "sd.linear_i8.epilogue." SD_CUDA_NAME ".v1";
|
||||||
desc.priority = 240;
|
desc.priority = 240;
|
||||||
desc.supports = nullptr;
|
desc.supports = nullptr;
|
||||||
desc.match = sd_i8_cuda::match;
|
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));
|
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.priority = 100;
|
||||||
desc.match = nullptr;
|
desc.match = nullptr;
|
||||||
desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
|
desc.name = "sd.quantize_i8_convrot." SD_CUDA_NAME ".v1";
|
||||||
desc.custom_op = sd_ggml_convrot_op();
|
desc.custom_op = sd_ggml_convrot_op();
|
||||||
desc.supports = sd_i8_cuda::convrot_supports;
|
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));
|
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));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
}
|
}
|
||||||
@@ -257,7 +257,7 @@ namespace SD_CUDA_NAMESPACE {
|
|||||||
if (j >= 2) {
|
if (j >= 2) {
|
||||||
const auto* mat = ggml_graph_node(graph, j - 2);
|
const auto* mat = ggml_graph_node(graph, j - 2);
|
||||||
const auto* view = ggml_graph_node(graph, j - 1);
|
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 &&
|
view->op == GGML_OP_RESHAPE && view->src[0] == mat && norm->src[0] == view &&
|
||||||
ggml_ext_graph_delay(graph, j - 2, 3, i))
|
ggml_ext_graph_delay(graph, j - 2, 3, i))
|
||||||
break;
|
break;
|
||||||
|
|||||||
@@ -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))
|
if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node))
|
||||||
return false;
|
return false;
|
||||||
*op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
|
*op = {node->src[0], node->src[1], nullptr, nullptr, nullptr, 0, false};
|
||||||
if (node->op == GGML_OP_MUL_MAT) {
|
if (ggml_ext_op_get(node) == sd_ggml_i8_op()) {
|
||||||
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()) {
|
|
||||||
I8Attributes attr;
|
I8Attributes attr;
|
||||||
std::memcpy(&attr, ggml_ext_op_attributes(node), sizeof(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)))
|
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) {
|
bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group) {
|
||||||
if (!node || !node->src[0] || !ggml_is_contiguous(node))
|
if (!node || !node->src[0] || !ggml_is_contiguous(node))
|
||||||
return false;
|
return false;
|
||||||
if (node->op == GGML_OP_QUANTIZE_I8_CONVROT)
|
if (ggml_ext_op_get(node) == sd_ggml_convrot_op())
|
||||||
*group = node->op_params[0];
|
|
||||||
else if (ggml_ext_op_get(node) == sd_ggml_convrot_op())
|
|
||||||
std::memcpy(group, ggml_ext_op_attributes(node), sizeof(int32_t));
|
std::memcpy(group, ggml_ext_op_attributes(node), sizeof(int32_t));
|
||||||
else
|
else
|
||||||
return false;
|
return false;
|
||||||
ggml_ext_tensor_desc output;
|
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));
|
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};
|
I8Attributes attr = {{logical->ne[1], logical->ne[2], logical->ne[3]}, group, logical != x ? I8_PACKED : 0};
|
||||||
ggml_tensor* inputs[5] = {w, x};
|
ggml_tensor* inputs[5] = {w, x};
|
||||||
int count = 2;
|
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);
|
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) {
|
ggml_tensor* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* x, int group) {
|
||||||
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) {
|
|
||||||
int32_t attr = group;
|
int32_t attr = group;
|
||||||
return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
|
return ggml_ext_op_create(ctx, sd_ggml_convrot_op(), &x, 1, &attr);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -32,9 +32,8 @@ ggml_ext_op_t sd_ggml_i8_op();
|
|||||||
ggml_ext_op_t sd_ggml_convrot_op();
|
ggml_ext_op_t sd_ggml_convrot_op();
|
||||||
bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation);
|
bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* operation);
|
||||||
bool sd_ggml_convrot_describe(const ggml_tensor* node, int* group_size);
|
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* 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* 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* ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
|
||||||
ggml_tensor* sd_ggml_quantize_i8_convrot(ggml_context* ctx, ggml_tensor* input, int group_size);
|
|
||||||
#endif
|
#endif
|
||||||
|
|
||||||
#endif
|
#endif
|
||||||
@@ -1,27 +1,69 @@
|
|||||||
find_package(Vulkan COMPONENTS glslc REQUIRED)
|
find_package(Vulkan COMPONENTS glslc REQUIRED)
|
||||||
set(SD_VULKAN_GENERATED "${CMAKE_CURRENT_BINARY_DIR}/shaders")
|
set(SD_VULKAN_GENERATED "${CMAKE_CURRENT_BINARY_DIR}/shaders")
|
||||||
file(MAKE_DIRECTORY "${SD_VULKAN_GENERATED}")
|
file(MAKE_DIRECTORY "${SD_VULKAN_GENERATED}")
|
||||||
foreach(NAME rope_apply rms_rope_apply rms_rope_cast)
|
|
||||||
set(SHADER "${NAME}")
|
function(sd_vulkan_feature NAME EXTENSION)
|
||||||
set(FLAGS)
|
set(SOURCE "${SD_VULKAN_GENERATED}/feature_${NAME}.comp")
|
||||||
if(NAME STREQUAL "rms_rope_cast")
|
set(EXTENSIONS "#extension ${EXTENSION} : require")
|
||||||
set(SHADER rms_rope_apply)
|
if(NAME STREQUAL "SD_VK_CM2")
|
||||||
set(FLAGS -DCAST_F16)
|
string(PREPEND EXTENSIONS "#extension GL_KHR_cooperative_matrix : require\n")
|
||||||
endif()
|
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(SOURCE "${CMAKE_CURRENT_SOURCE_DIR}/shaders/${SHADER}.comp")
|
||||||
set(SPV "${SD_VULKAN_GENERATED}/${NAME}.spv")
|
set(SPV "${SD_VULKAN_GENERATED}/${NAME}.spv")
|
||||||
set(HEADER "${SD_VULKAN_GENERATED}/${NAME}.h")
|
set(ENV vulkan1.2)
|
||||||
# Keep native RoPE rounding: spirv-opt changes the subgroup reduction results.
|
if(NAME MATCHES "_cm2$")
|
||||||
|
set(ENV vulkan1.3)
|
||||||
|
endif()
|
||||||
add_custom_command(OUTPUT "${HEADER}"
|
add_custom_command(OUTPUT "${HEADER}"
|
||||||
COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=vulkan1.2 ${FLAGS} "${SOURCE}" -o "${SPV}"
|
COMMAND "${Vulkan_GLSLC_EXECUTABLE}" --target-env=${ENV} ${ARGN} "${SOURCE}" -o "${SPV}"
|
||||||
COMMAND "${CMAKE_COMMAND}" "-DINPUT=${SPV}" "-DOUTPUT=${HEADER}" "-DNAME=${NAME}"
|
COMMAND "${CMAKE_COMMAND}" "-DINPUT=${SPV}" "-DOUTPUT=${HEADER}" "-DNAME=${NAME}"
|
||||||
-P "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake"
|
-P "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake"
|
||||||
DEPENDS "${SOURCE}" "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake"
|
DEPENDS "${SOURCE}" "${PROJECT_SOURCE_DIR}/cmake/embed_spirv.cmake"
|
||||||
BYPRODUCTS "${SPV}"
|
BYPRODUCTS "${SPV}" VERBATIM)
|
||||||
VERBATIM)
|
else()
|
||||||
list(APPEND SD_VULKAN_HEADERS "${HEADER}")
|
file(WRITE "${HEADER}" "#include <cstdint>\nnamespace sd_vulkan_shaders { static constexpr uint32_t ${NAME}[] = {0}; }\n")
|
||||||
endforeach()
|
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_custom_target(sd-vulkan-shaders DEPENDS ${SD_VULKAN_HEADERS})
|
||||||
add_dependencies(${SD_LIB} sd-vulkan-shaders)
|
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}")
|
target_include_directories(${SD_LIB} PRIVATE "${SD_VULKAN_GENERATED}")
|
||||||
@@ -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));
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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();
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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;
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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);
|
||||||
|
}
|
||||||
|
}
|
||||||
@@ -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();
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in new issue
Block a user