refactor: preserve INT8 APIs through backend extensions

This commit is contained in:
leejet committed 2026-11-01 01:55:51 +08:00
1 parent 051b49751e
commit f0204a7eaa
20 files changed
+1573 -102

No files matched your search

+1 -1
View File
@@ -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}")
+1
View File
@@ -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
View File
@@ -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.
-8
View File
@@ -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));
} }
} }
+10 -6
View File
@@ -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));
} }
+2 -6
View File
@@ -449,7 +449,8 @@ void sd_ggml_register_cuda_i8_gated_residual(ggml_backend_dev_t device) {
ggml_ext_kernel_desc desc = {}; 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));
} }
+10 -5
View File
@@ -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
View File
@@ -462,8 +462,8 @@ namespace SD_CUDA_NAMESPACE {
p->scales = reserve(total, op.input->type == GGML_TYPE_F32 ? size_t(p->rows) * sizeof(float) : 0); 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));
} }
} }
+1 -1
View File
@@ -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;
+13 -30
View File
@@ -161,14 +161,7 @@ bool sd_ggml_i8_describe(const ggml_tensor* node, SDI8Operation* op) {
if (!node || node->type != GGML_TYPE_F32 || !ggml_is_contiguous(node)) 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);
} }
+2 -3
View File
@@ -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
+56 -14
View File
@@ -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}")
+771
View File
@@ -0,0 +1,771 @@
#include <algorithm>
#include <array>
#include <climits>
#include <cstring>
#include <vector>
#include "core/ggml_extensions.h"
#include "ggml-ext-vulkan.h"
#include "modulate_convrot.h"
#include "mul_mat_i8_residual_cm2.h"
#include "mul_mat_i8_swiglu_split_cm2.h"
#include "mul_mat_i8_tensorwise.h"
#include "mul_mat_i8_tensorwise_cm1.h"
#include "mul_mat_i8_tensorwise_cm2.h"
#include "quantize_i8_convrot.h"
#include "quantize_i8_convrot_register.h"
#include "quantize_i8_convrot_subgroup.h"
#include "rms_modulate_convrot.h"
#include "sigmoid_quantize_i8_convrot_6144.h"
#include "swiglu_quantize_i8_convrot.h"
#include "swiglu_quantize_i8_convrot_16384.h"
namespace sd_vulkan_i8 {
struct Device {
const ggml_ext_vulkan_api* api;
ggml_ext_vulkan_device_info info;
ggml_ext_vulkan_compute_info compute;
bool subgroup, cm2, cm2_small;
uint32_t cm_m, cm_n, cm_k;
};
static bool device_info(ggml_backend_dev_t dev, Device& d) {
auto reg = ggml_backend_dev_backend_reg(dev);
auto query = reinterpret_cast<ggml_ext_vulkan_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_v1"));
auto compute = reinterpret_cast<ggml_ext_vulkan_compute_api_fn>(ggml_backend_reg_get_proc_address(reg, "ggml_ext_vulkan_compute_v1"));
d = {};
if (!query || !compute)
return false;
d.api = query();
auto service = compute();
if (!d.api->device_info(dev, &d.info) || !service->properties(dev, &d.compute))
return false;
d.subgroup = d.info.subgroup_size == 32 && d.info.full_subgroups && d.info.subgroup_shuffle && d.compute.subgroup_arithmetic && d.info.max_group_invocations >= 256;
ggml_ext_vulkan_matrix_properties prop;
for (uint32_t i = 0; service->matrix_properties(dev, i, &prop); ++i) {
if (prop.a != GGML_TYPE_I8 || prop.b != GGML_TYPE_I8 || prop.accumulator != GGML_TYPE_I32 || prop.result != GGML_TYPE_I32 || prop.saturating || !prop.m || !prop.n || !prop.k)
continue;
if (!prop.flexible && !prop.workgroup_scope && !d.cm_m && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm1) > 4) {
d.cm_m = prop.m;
d.cm_n = prop.n;
d.cm_k = prop.k;
}
if (prop.flexible && prop.workgroup_scope && prop.invocations == 128 && 128 % prop.m == 0 && 64 % prop.n == 0 && 64 % prop.k == 0 && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise_cm2) > 4) {
d.cm2 = true;
d.cm2_small |= 64 % prop.m == 0 && 32 % prop.n == 0;
}
}
return d.info.integer_dot_product && sizeof(sd_vulkan_shaders::mul_mat_i8_tensorwise) > 4;
}
static const ggml_tensor* node(const ggml_cgraph* graph, int i) {
return ggml_graph_node(const_cast<ggml_cgraph*>(graph), i);
}
static int count_nodes(const ggml_cgraph* graph) {
return ggml_graph_n_nodes(const_cast<ggml_cgraph*>(graph));
}
static int param_i32(const ggml_tensor* node, int i) {
return node->op_params[i];
}
static float param_f32(const ggml_tensor* node, int i) {
float value;
std::memcpy(&value, node->op_params + i, 4);
return value;
}
static bool is_convrot(const ggml_tensor* node) {
int group;
return sd_ggml_convrot_describe(node, &group);
}
static bool is_mat(const ggml_tensor* node) {
SDI8Operation op;
return sd_ggml_i8_describe(node, &op);
}
static bool aligned(const ggml_tensor* t, size_t alignment) {
size_t offset = 0;
while (!t->data && t->view_src) {
offset += t->view_offs;
t = t->view_src;
}
return (reinterpret_cast<uintptr_t>(t->data) + offset) % alignment == 0;
}
static bool bounded(const Device& d, const ggml_tensor* t, bool descriptor_aligned = true) {
return t && ggml_nbytes(t) <= d.info.max_buffer_range - (d.info.buffer_alignment - 1) &&
aligned(t, descriptor_aligned ? d.info.buffer_alignment : 4);
}
static bool can_modulate(const Device& d, int64_t k, bool affine) {
return (k == 4096 || k == 6144 || (affine && k == 5376)) && d.info.max_group_invocations >= 512 && d.info.max_group_size[0] >= 512 && d.compute.shared_memory >= (k + 1024) * 4;
}
static bool can_quantize(const Device& d, const ggml_tensor* dst) {
int group;
if (!sd_ggml_convrot_describe(dst, &group) || group != 256 || !bounded(d, dst) || !bounded(d, dst->src[0]))
return false;
const auto* x = dst->src[0];
return x->ne[0] <= UINT32_MAX && ggml_nrows(x) <= UINT32_MAX && d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 256;
}
static bool can_matmul(const Device& d, const ggml_tensor* dst) {
SDI8Operation op;
if (!sd_ggml_i8_describe(dst, &op) || !op.packed || op.group_size != 256 || !op.weight_scale || !bounded(d, dst))
return false;
for (auto* t : {op.weight, op.input, op.weight_scale, op.bias})
if (t && !bounded(d, t))
return false;
const uint64_t cols = op.weight->ne[1], rows = ggml_nrows(dst);
return op.weight->ne[0] % 256 == 0 && op.weight->ne[0] <= UINT32_MAX && cols <= UINT32_MAX && rows <= UINT32_MAX &&
(cols + 15) / 16 <= d.info.max_group_count[0] && (rows + 15) / 16 <= d.info.max_group_count[1] &&
d.info.max_group_invocations >= 256 && d.info.max_group_size[0] >= 16 && d.info.max_group_size[1] >= 16;
}
static bool can_fuse_sigmoid_convrot(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 2 >= count_nodes(graph) || !device.subgroup) {
return false;
}
const ggml_tensor* sigmoid = node(graph, i);
const ggml_tensor* mul = node(graph, i + 1);
const ggml_tensor* quant = node(graph, i + 2);
if (sigmoid->op != GGML_OP_UNARY || ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID ||
mul->op != GGML_OP_MUL || mul->src[1] != sigmoid ||
!is_convrot(quant) || quant->src[0] != mul ||
sigmoid->ne[0] != 6144 || !ggml_are_same_shape(sigmoid, mul->src[0]) ||
!can_quantize(device, quant)) {
return false;
}
for (const ggml_tensor* tensor : {sigmoid->src[0], mul->src[0]}) {
if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) ||
tensor->view_offs % device.info.buffer_alignment != 0) {
return false;
}
}
return true;
}
static bool can_fuse_swiglu_convrot(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 1 >= count_nodes(graph) || !device.subgroup) {
return false;
}
const ggml_tensor* glu = node(graph, i);
const ggml_tensor* quant = node(graph, i + 1);
if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
glu->src[1] == nullptr || param_i32(glu, 1) != 0 ||
!is_convrot(quant) || quant->src[0] != glu ||
!can_quantize(device, quant) || !ggml_is_contiguous(glu->src[0]) || !ggml_is_contiguous(glu->src[1])) {
return false;
}
for (const ggml_tensor* tensor : {glu->src[0], glu->src[1]}) {
if (tensor->type != GGML_TYPE_F32 || tensor->view_offs % device.info.buffer_alignment != 0) {
return false;
}
}
return true;
}
static int match_modulate_convrot(const Device& device, const ggml_cgraph* graph, int quant_idx, int* indices) {
if (quant_idx < 1 || quant_idx >= count_nodes(graph)) {
return 0;
}
const ggml_tensor* quant = node(graph, quant_idx);
if (!is_convrot(quant) || !can_quantize(device, quant)) {
return 0;
}
const ggml_tensor* nodes[4]{};
const ggml_tensor* mod = nullptr;
const bool affine = quant->src[0]->op == GGML_OP_ADD;
const int count = affine ? 4 : 3;
if (affine) {
const ggml_tensor* shift = quant->src[0];
const ggml_tensor* add = shift->src[0];
if (add->op != GGML_OP_ADD || add->src[1]->op != GGML_OP_MUL || add->src[1]->src[0] != add->src[0]) {
return 0;
}
nodes[0] = add->src[1];
nodes[1] = add;
nodes[2] = shift;
nodes[3] = quant;
mod = nodes[0]->src[1];
const ggml_tensor* vector = shift->src[1];
if (vector->type != GGML_TYPE_F32 || !ggml_is_contiguous(vector) ||
vector->ne[0] != add->ne[0] || ggml_nelements(vector) != add->ne[0]) {
return 0;
}
} else {
const ggml_tensor* mul = quant->src[0];
if (mul->op != GGML_OP_MUL || mul->src[0]->op != GGML_OP_NORM) {
return 0;
}
nodes[0] = mul->src[0];
nodes[1] = mul;
nodes[2] = quant;
mod = mul->src[1];
}
const ggml_tensor* x = nodes[0]->src[0];
if ((x->ne[0] != 4096 && x->ne[0] != 6144 && x->ne[0] != 5376) ||
!can_modulate(device, x->ne[0], affine) ||
x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || mod->type != GGML_TYPE_F32 ||
!ggml_is_contiguous(mod) || mod->ne[0] != x->ne[0] || ggml_nelements(mod) != x->ne[0]) {
return 0;
}
for (int j = 0; j < count; ++j) {
if (j < count - 1 && (nodes[j]->type != GGML_TYPE_F32 || !ggml_is_contiguous(nodes[j]) || !ggml_are_same_shape(x, nodes[j]))) {
return 0;
}
}
int next = count - 1;
for (int j = quant_idx; j >= std::max(0, quant_idx - 64) && next >= 0; --j) {
if (node(graph, j) == nodes[next]) {
indices[next--] = j;
}
}
return next < 0 ? count : 0;
}
static int match_rms_modulate_convrot(const Device& device, const ggml_cgraph* graph, int q, int* indices) {
if (!can_modulate(device, 5376, true) || match_modulate_convrot(device, graph, q, indices + 2) != 4) {
return 0;
}
const auto weighted = node(graph, indices[2])->src[0];
if (weighted->op != GGML_OP_MUL || weighted->src[0]->op != GGML_OP_RMS_NORM || weighted->ne[0] != 5376) {
return 0;
}
const auto norm = weighted->src[0];
const auto x = norm->src[0];
const auto weight = weighted->src[1];
if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || weight->type != GGML_TYPE_F32 ||
!ggml_is_contiguous(weight) || ggml_nelements(weight) != 5376 || weight->ne[0] != 5376) {
return 0;
}
int next = 1;
for (int j = indices[2] - 1; j >= std::max(0, indices[2] - 64) && next >= 0; --j) {
if (node(graph, j) == (next == 1 ? weighted : norm)) {
indices[next--] = j;
}
}
return next < 0 ? 6 : 0;
}
static int can_fuse_modulate_convrot(const Device& device, const ggml_cgraph* graph, int i) {
const ggml_op op = node(graph, i)->op;
if (op == GGML_OP_RMS_NORM) {
int indices[6];
if (match_rms_modulate_convrot(device, graph, i + 5, indices) != 6) {
return 0;
}
for (int j = 0; j < 6; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return 6;
}
if (op != GGML_OP_NORM && op != GGML_OP_MUL) {
return 0;
}
const int count = op == GGML_OP_NORM ? 3 : 4;
int indices[4];
if (match_modulate_convrot(device, graph, i + count - 1, indices) != count) {
return 0;
}
for (int j = 0; j < count; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return count;
}
static bool match_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i, int* indices) {
if (i < 2 || i >= count_nodes(graph) || !device.cm2_small) {
return false;
}
const ggml_tensor* glu = node(graph, i);
if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU ||
param_i32(glu, 1) != 0 || glu->src[1] == nullptr || glu->ne[0] % 128 != 0 ||
glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) ||
glu->view_offs % device.info.buffer_alignment != 0) {
return false;
}
const ggml_tensor* gate = glu->src[0];
const ggml_tensor* up = glu->src[1];
if (gate == up || gate->src[1] != up->src[1] || !ggml_are_same_shape(gate, up)) {
return false;
}
for (auto mat : {gate, up}) {
if (!is_mat(mat) || mat->view_src != nullptr || mat->src[3] != nullptr ||
!can_matmul(device, mat)) {
return false;
}
}
int found = 0;
for (int j = std::max(0, i - 32); j < i; ++j) {
if (node(graph, j) == gate || node(graph, j) == up) {
indices[found++] = j;
}
}
if (found != 2) {
return false;
}
indices[2] = i;
return true;
}
static bool can_fuse_i8_swiglu_split(const Device& device, const ggml_cgraph* graph, int i) {
int indices[3];
return match_i8_swiglu_split(device, graph, i + 2, indices) && indices[0] == i && indices[1] == i + 1;
}
static bool can_fuse_i8_swiglu(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 1 >= count_nodes(graph) || !device.cm2_small) {
return false;
}
const ggml_tensor* mat = node(graph, i);
const ggml_tensor* glu = node(graph, i + 1);
return is_mat(mat) && can_matmul(device, mat) &&
mat->view_src == nullptr && mat->src[3] == nullptr && mat->ne[0] % 256 == 0 &&
glu->op == GGML_OP_GLU && ggml_get_glu_op(glu) == GGML_GLU_OP_SWIGLU &&
param_i32(glu, 1) == 0 && glu->src[0] == mat && glu->src[1] == nullptr &&
glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) &&
glu->view_offs % device.info.buffer_alignment == 0;
}
static int match_i8_residual(const Device& device, const ggml_cgraph* graph, int add_idx, int* indices) {
if (add_idx < 2 || add_idx >= count_nodes(graph) || !device.cm2) {
return 0;
}
const ggml_tensor* dst = node(graph, add_idx);
if (dst->op != GGML_OP_ADD) {
return 0;
}
const ggml_tensor* mul = dst->src[0]->op == GGML_OP_MUL ? dst->src[0] : dst->src[1];
if (mul->op != GGML_OP_MUL) {
return 0;
}
const ggml_tensor* scaled_mat = mul->src[0];
const bool scaled = scaled_mat->op == GGML_OP_SCALE;
const ggml_tensor* mat = scaled ? scaled_mat->src[0] : scaled_mat;
if (!is_mat(mat) || !can_matmul(device, mat) || mat->src[3] != nullptr ||
(scaled && param_f32(scaled_mat, 1) != 0.0f)) {
return 0;
}
const int count = scaled ? 4 : 3;
const ggml_tensor* nodes[4] = {mat, scaled ? scaled_mat : mul, scaled ? mul : dst, dst};
const ggml_tensor* gate = mul->src[1];
const ggml_tensor* residual = dst->src[0] == mul ? dst->src[1] : dst->src[0];
if (!ggml_are_same_shape(mat, dst) || !ggml_are_same_shape(residual, dst) ||
gate->ne[0] != mat->ne[0] || ggml_nelements(gate) != mat->ne[0]) {
return 0;
}
for (const ggml_tensor* tensor : {mat, scaled_mat, mul, gate, residual, dst}) {
if (tensor->type != GGML_TYPE_F32 || !ggml_is_contiguous(tensor) ||
tensor->view_offs % device.info.buffer_alignment != 0) {
return 0;
}
}
int next = count - 1;
for (int j = add_idx; j >= std::max(0, add_idx - 32) && next >= 0; --j) {
if (node(graph, j) == nodes[next]) {
indices[next--] = j;
}
}
return next < 0 ? count : 0;
}
static int can_fuse_i8_residual(const Device& device, const ggml_cgraph* graph, int i) {
if (i + 2 >= count_nodes(graph) || !is_mat(node(graph, i))) {
return 0;
}
const int count = node(graph, i + 1)->op == GGML_OP_SCALE ? 4 : 3;
int indices[4];
if (match_i8_residual(device, graph, i + count - 1, indices) != count) {
return 0;
}
for (int j = 0; j < count; ++j) {
if (indices[j] != i + j) {
return 0;
}
}
return count;
}
static int match_i8_row_concat(const Device& device, const ggml_cgraph* graph, int end, int* indices) {
const ggml_tensor* concat = node(graph, end);
if (concat->op != GGML_OP_CONCAT || param_i32(concat, 0) != 1 ||
concat->type != GGML_TYPE_F32 || concat->ne[2] != 1 || concat->ne[3] != 1 ||
!ggml_is_contiguous(concat) || !ggml_is_contiguous(concat->src[0]) ||
concat->view_offs % device.info.buffer_alignment != 0 ||
ggml_nbytes(concat->src[0]) % device.info.buffer_alignment != 0) {
return 0;
}
const ggml_tensor* tail = concat->src[1];
int count = 0;
for (int j = end - 1; j >= std::max(0, end - 32); --j) {
if (node(graph, j) != tail) {
continue;
}
if (is_mat(tail) && can_matmul(device, tail)) {
indices[0] = j;
count = 1;
} else {
count = match_i8_residual(device, graph, j, indices);
}
break;
}
if (!count) {
return 0;
}
indices[count++] = end;
return count;
}
static int can_fuse_i8_row_concat(const Device& device, const ggml_cgraph* graph, int i) {
if (!is_mat(node(graph, i))) {
return 0;
}
for (int end = i + 1; end < std::min(i + 5, count_nodes(graph)); ++end) {
int indices[5];
const int count = match_i8_row_concat(device, graph, end, indices);
if (count != end - i + 1) {
continue;
}
bool consecutive = true;
for (int j = 0; j < count; ++j) {
consecutive &= indices[j] == i + j;
}
if (consecutive) {
return count;
}
}
return 0;
}
struct QuantParameters {
uint32_t k, rows;
};
struct ModulateParameters {
uint32_t rows, x_offset, m_offset, h_offset;
float eps;
uint32_t k, w_offset;
};
struct MatParameters {
uint32_t n, rows, k, has_bias, row_offset;
float output_scale;
uint32_t group_cols;
};
static_assert(sizeof(QuantParameters) == 8 && sizeof(ModulateParameters) == 28 && sizeof(MatParameters) == 28, "INT8 push constant layout mismatch");
template <size_t N>
static ggml_ext_vulkan_pipeline shader(const char* name, const uint32_t (&spirv)[N], uint32_t bindings, uint32_t constants) {
return {name, spirv, sizeof(spirv), bindings, constants, nullptr, 0, 0, false};
}
#define SD_VK_PIPELINE(NAME, BINDINGS, PARAMS) shader("sd." #NAME ".vulkan", sd_vulkan_shaders::NAME, BINDINGS, sizeof(PARAMS))
static ggml_ext_vulkan_binding range(ggml_ext_vulkan_binding binding, size_t offset, size_t size) {
binding.offset += offset;
binding.size = size;
return binding;
}
static bool overlap(const ggml_ext_vulkan_buffer_info& a, const ggml_ext_vulkan_buffer_info& b) {
return a.storage == b.storage && a.offset < b.offset + b.size && b.offset < a.offset + a.size;
}
static bool output_binding(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const std::vector<ggml_ext_vulkan_binding>& inputs, ggml_ext_vulkan_binding& binding, bool& alias) {
binding = {output};
alias = false;
ggml_ext_vulkan_buffer_info dst, src;
if (!d.api->buffer_info(e->backend, &binding, &dst))
return false;
for (const auto& input : inputs) {
if (!d.api->buffer_info(e->backend, &input, &src))
return false;
alias |= overlap(dst, src);
}
if (alias)
binding = {nullptr, e->workspace, 0, ggml_nbytes(output)};
return true;
}
static ggml_status finish(const ggml_ext_execution_context* e, const Device& d, const ggml_tensor* output, const ggml_ext_vulkan_binding& binding, bool alias, ggml_status status) {
if (status == GGML_STATUS_SUCCESS && alias) {
ggml_ext_vulkan_binding dst = {output};
return d.api->copy(e->backend, &binding, &dst);
}
return status;
}
static ggml_status quantize(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
const auto* output = nodes[count - 1];
const auto* first = nodes[0];
const bool modulate = count > 1 && (first->op == GGML_OP_NORM || first->op == GGML_OP_MUL || first->op == GGML_OP_RMS_NORM);
const bool rms = modulate && first->op == GGML_OP_RMS_NORM;
const bool affine = modulate && (rms || first->op == GGML_OP_MUL);
const bool sigmoid = count > 1 && first->op == GGML_OP_UNARY;
const auto* x = count == 1 ? output->src[0] : first->src[0];
const auto* up = count == 1 || modulate ? nullptr : sigmoid ? nodes[1]->src[0]
: first->src[1];
const int extra = rms ? 2 : 0;
const auto* mod = modulate ? nodes[extra + (affine ? 0 : 1)]->src[1] : nullptr;
const auto* shift = modulate ? (affine ? nodes[extra + 2]->src[1] : mod) : nullptr;
const auto* weight = rms ? nodes[1]->src[1] : x;
std::vector<ggml_ext_vulkan_binding> inputs = {{x}};
if (up)
inputs.push_back({up});
if (modulate) {
inputs.push_back({mod});
inputs.push_back({shift});
if (rms)
inputs.push_back({weight});
}
ggml_ext_vulkan_binding dst;
bool alias;
if (!output_binding(e, d, output, inputs, dst, alias))
return GGML_STATUS_FAILED;
const uint32_t k = uint32_t(x->ne[0]), rows = uint32_t(ggml_nrows(x));
const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
auto packed = range(dst, 0, scale_offset), scales = range(dst, scale_offset, size_t(rows) * 4);
const uint32_t groups[] = {std::min(rows, d.info.max_group_count[0]), 1, 1};
ggml_status status;
if (modulate) {
std::vector<ggml_ext_vulkan_binding> bindings = {{x}, {mod}, {shift}, packed, scales};
if (rms)
bindings.push_back({weight});
ModulateParameters pc = {rows, 0, 0, 0, affine && !rms ? 0.f : param_f32(first, 0), k, 0};
uint32_t* offsets[] = {&pc.x_offset, &pc.m_offset, &pc.h_offset, nullptr, nullptr, &pc.w_offset};
for (size_t i = 0; i < bindings.size(); ++i)
if (offsets[i]) {
ggml_ext_vulkan_buffer_info info;
if (!d.api->buffer_info(e->backend, &bindings[i], &info))
return GGML_STATUS_FAILED;
*offsets[i] = info.shader_offset / 4;
}
auto pipeline = rms ? SD_VK_PIPELINE(rms_modulate_convrot, 6, ModulateParameters) : SD_VK_PIPELINE(modulate_convrot, 5, ModulateParameters);
const uint32_t constants[] = {k, uint32_t(affine)};
pipeline.constants = constants;
pipeline.constant_count = 2;
status = d.api->dispatch(e->backend, &pipeline, bindings.data(), &pc, groups);
} else {
ggml_ext_vulkan_binding bindings[] = {{x}, packed, scales, {up}};
auto pipeline = SD_VK_PIPELINE(quantize_i8_convrot, 3, QuantParameters);
bool subgroup = d.subgroup && k >= 1024;
uint32_t width = k;
if (subgroup) {
pipeline = SD_VK_PIPELINE(quantize_i8_convrot_subgroup, 3, QuantParameters);
if (!up && ((uint64_t(k) + 2047) / 2048) * 8 <= 64) {
pipeline = SD_VK_PIPELINE(quantize_i8_convrot_register, 3, QuantParameters);
pipeline.constants = &width;
pipeline.constant_count = 1;
}
}
if (up) {
subgroup = true;
pipeline = k == 16384 ? SD_VK_PIPELINE(swiglu_quantize_i8_convrot_16384, 4, QuantParameters) : SD_VK_PIPELINE(swiglu_quantize_i8_convrot, 4, QuantParameters);
if (sigmoid)
pipeline = SD_VK_PIPELINE(sigmoid_quantize_i8_convrot_6144, 4, QuantParameters);
}
if (subgroup) {
pipeline.full_subgroups = d.info.full_subgroups;
pipeline.required_subgroup_size = d.info.subgroup_size_control ? 32 : 0;
}
QuantParameters pc = {k, rows};
status = d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups);
}
return finish(e, d, output, dst, alias, status);
}
static uint32_t group_columns(uint32_t cols, uint32_t k, uint32_t tile_cols, bool swiglu) {
if (!cols || !k || !tile_cols)
return 1;
const uint64_t col_tiles = (uint64_t(cols) + tile_cols - 1) / tile_cols;
const uint64_t weight_tile_bytes = uint64_t(tile_cols) * k * (swiglu ? 2 : 1);
const uint64_t growth = std::clamp(uint64_t(k) / cols, uint64_t(1), uint64_t(2));
return uint32_t(std::clamp(uint64_t(8) * 1024 * 1024 * growth / weight_tile_bytes, uint64_t(1), std::min(col_tiles, uint64_t(32))));
}
static ggml_status multiply(const ggml_ext_execution_context* e, const Device& d, ggml_tensor* const* nodes, int count) {
const auto* output = nodes[count - 1];
const auto* concat = output->op == GGML_OP_CONCAT ? output : nullptr;
const auto* tail = concat ? concat->src[1] : output;
const auto* glu = tail->op == GGML_OP_GLU ? tail : nullptr;
const auto* mat = glu ? glu->src[0] : tail;
const ggml_tensor *gate = nullptr, *residual = nullptr;
float output_scale = 1.f;
if (tail->op == GGML_OP_ADD) {
const auto* mul = tail->src[0]->op == GGML_OP_MUL ? tail->src[0] : tail->src[1];
gate = mul->src[1];
residual = tail->src[0] == mul ? tail->src[1] : tail->src[0];
mat = mul->src[0];
if (mat->op == GGML_OP_SCALE) {
output_scale = param_f32(mat, 0);
mat = mat->src[0];
}
}
const auto* up_mat = glu ? glu->src[1] : nullptr;
SDI8Operation op, up;
if (!sd_ggml_i8_describe(mat, &op) || (up_mat && !sd_ggml_i8_describe(up_mat, &up)))
return GGML_STATUS_FAILED;
const uint32_t k = uint32_t(op.weight->ne[0]), n = uint32_t(op.weight->ne[1]), rows = uint32_t(ggml_nrows(mat));
const size_t scale_offset = size_t(k) * ((uint64_t(rows) + 3) / 4 * 4);
std::vector<ggml_ext_vulkan_binding> inputs = {{op.weight}, {op.input}, {op.weight_scale}, {op.bias ? op.bias : op.weight_scale}};
if (gate) {
inputs.push_back({gate});
inputs.push_back({residual});
}
if (up_mat) {
inputs.push_back({up.weight});
inputs.push_back({up.weight_scale});
}
if (concat && ggml_nbytes(concat->src[0]))
inputs.push_back({concat->src[0]});
ggml_ext_vulkan_binding dst;
bool alias;
if (!output_binding(e, d, output, inputs, dst, alias))
return GGML_STATUS_FAILED;
auto destination = concat ? range(dst, ggml_nbytes(concat->src[0]), ggml_nbytes(tail)) : dst;
ggml_ext_vulkan_binding bindings[] = {{op.weight}, {op.input, nullptr, 0, scale_offset}, {op.weight_scale}, {op.input, nullptr, scale_offset, size_t(rows) * 4}, {op.bias ? op.bias : op.weight_scale}, destination, {up_mat ? up.weight : gate}, {up_mat ? up.weight_scale : residual}};
auto dispatch = [&](ggml_ext_vulkan_pipeline pipeline, uint32_t tile_cols, uint32_t tile_rows,
std::initializer_list<uint32_t> constants, uint32_t row_offset, uint32_t dispatch_rows) {
if (!dispatch_rows)
return GGML_STATUS_SUCCESS;
const uint32_t cols = glu ? uint32_t(glu->ne[0]) : n;
MatParameters pc = {n * (up_mat ? 2u : 1u), rows, k, op.bias ? 1u : 0u, row_offset, output_scale, group_columns(cols, k, tile_cols, glu != nullptr)};
pipeline.constants = constants.begin();
pipeline.constant_count = uint32_t(constants.size());
const uint32_t groups[] = {(cols + tile_cols - 1) / tile_cols, (dispatch_rows + tile_rows - 1) / tile_rows, 1};
return d.api->dispatch(e->backend, &pipeline, bindings, &pc, groups);
};
ggml_status status = GGML_STATUS_SUCCESS;
if (d.cm2) {
const bool large_k = k % 4096 == 0 || (rows >= n && k >= n);
const uint32_t full_rows = (glu || n % 128 == 0) ? rows / 128 * 128 : 0;
auto pipeline = residual ? SD_VK_PIPELINE(mul_mat_i8_residual_cm2, 8, MatParameters) : up_mat ? SD_VK_PIPELINE(mul_mat_i8_swiglu_split_cm2, 8, MatParameters)
: SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm2, 6, MatParameters);
const uint32_t tile_cols = glu ? 64 : 128;
status = dispatch(pipeline, tile_cols, 128, {128, 1, glu ? 1u : 0u, tile_cols, glu || large_k ? 128u : 64u}, 0, full_rows);
if (status == GGML_STATUS_SUCCESS) {
const bool small = d.cm2_small && (glu || n % 128 == 0) && rows - full_rows <= 32;
const uint32_t tail_cols = small || glu ? 64 : 128, tail_rows = small ? 32 : glu ? 128
: 64;
status = dispatch(pipeline, tail_cols, tail_rows, {tail_rows, small ? 2u : 0u, glu ? 1u : 0u, tail_cols, small ? 256u : 64u}, full_rows, rows - full_rows);
}
} else {
uint32_t coop_rows = 0;
if (d.cm_m && k % d.cm_k == 0 && n % (4 * d.cm_n) == 0) {
const uint32_t tile_rows = 4 * d.cm_m;
coop_rows = rows / tile_rows * tile_rows;
auto pipeline = SD_VK_PIPELINE(mul_mat_i8_tensorwise_cm1, 6, MatParameters);
pipeline.full_subgroups = d.info.full_subgroups;
pipeline.required_subgroup_size = d.info.subgroup_size_control ? d.info.subgroup_size : 0;
status = dispatch(pipeline, 4 * d.cm_n, tile_rows, {d.info.subgroup_size, d.cm_m, d.cm_n, d.cm_k}, 0, coop_rows);
}
if (status == GGML_STATUS_SUCCESS)
status = dispatch(SD_VK_PIPELINE(mul_mat_i8_tensorwise, 6, MatParameters), 16, 16, {}, coop_rows, rows - coop_rows);
}
if (status == GGML_STATUS_SUCCESS && concat && ggml_nbytes(concat->src[0])) {
ggml_ext_vulkan_binding src = {concat->src[0]}, prefix = range(dst, 0, ggml_nbytes(concat->src[0]));
status = d.api->copy(e->backend, &src, &prefix);
}
return finish(e, d, output, dst, alias, status);
}
static bool quant_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
Device d;
return device_info(dev, d) && can_quantize(d, n);
}
static bool mat_supports(ggml_backend_dev_t dev, const ggml_tensor* n, void*) {
Device d;
return device_info(dev, d) && can_matmul(d, n);
}
static int match(ggml_backend_dev_t dev, const ggml_cgraph* graph, int first, void*) {
Device d;
if (!device_info(dev, d) || d.info.fusion_disabled)
return 0;
if (int count = can_fuse_modulate_convrot(d, graph, first))
return count;
if (int count = can_fuse_i8_row_concat(d, graph, first))
return count;
if (int count = can_fuse_i8_residual(d, graph, first))
return count;
if (can_fuse_i8_swiglu_split(d, graph, first))
return 3;
if (can_fuse_i8_swiglu(d, graph, first))
return 2;
if (can_fuse_sigmoid_convrot(d, graph, first))
return 3;
if (can_fuse_swiglu_convrot(d, graph, first))
return 2;
return 0;
}
static void optimize(ggml_backend_dev_t dev, ggml_cgraph* graph, void*) {
Device d;
if (!device_info(dev, d) || d.info.fusion_disabled || d.compute.graph_optimization_disabled)
return;
for (int last = 1; last < count_nodes(graph); ++last) {
int indices[6];
int count = match_rms_modulate_convrot(d, graph, last, indices);
if (!count)
count = match_modulate_convrot(d, graph, last, indices);
if (!count)
count = match_i8_row_concat(d, graph, last, indices);
if (!count)
count = match_i8_residual(d, graph, last, indices);
if (!count && match_i8_swiglu_split(d, graph, last, indices))
count = 3;
if (!count)
continue;
int first = indices[0];
for (int j = 1; j < count; ++j) {
if (!ggml_ext_graph_delay(graph, first, j, indices[j]))
break;
first = indices[j] - j;
}
}
}
static ggml_status prepare(ggml_backend_t, ggml_tensor* const* nodes, int count, void**, size_t* bytes, void*) {
*bytes = ggml_nbytes(nodes[count - 1]);
return GGML_STATUS_SUCCESS;
}
static void destroy(void*, void*) {}
static ggml_status execute(const ggml_ext_execution_context* e, ggml_tensor* const* nodes, int count, void*, void*) {
Device d;
if (!device_info(ggml_backend_get_device(e->backend), d))
return GGML_STATUS_FAILED;
return is_convrot(nodes[count - 1]) ? quantize(e, d, nodes, count) : multiply(e, d, nodes, count);
}
#undef SD_VK_PIPELINE
}
void sd_ggml_register_vulkan_i8(ggml_backend_dev_t device) {
using namespace sd_vulkan_i8;
Device info;
if (!device_info(device, info))
return;
ggml_ext_kernel_desc d = {};
d.abi_version = GGML_EXT_ABI_VERSION;
d.priority = 100;
d.prepare = prepare;
d.destroy = destroy;
d.execute = execute;
d.name = "sd.quantize_i8_convrot.vulkan.v1";
d.op = GGML_OP_CUSTOM;
d.custom_op = sd_ggml_convrot_op();
d.supports = quant_supports;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.name = "sd.linear_i8.vulkan.v1";
d.op = GGML_OP_CUSTOM;
d.custom_op = sd_ggml_i8_op();
d.supports = mat_supports;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.priority = 200;
d.supports = nullptr;
d.match = match;
d.optimize = optimize;
d.name = "sd.linear_i8.fused.vulkan.v1";
d.op = GGML_OP_CUSTOM;
d.custom_op = sd_ggml_i8_op();
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
d.optimize = nullptr;
d.custom_op = nullptr;
const struct {
ggml_op op;
const char* name;
} fusions[] = {
{GGML_OP_NORM, "sd.norm_convrot.vulkan.v1"},
{GGML_OP_RMS_NORM, "sd.rms_modulate_convrot.vulkan.v1"},
{GGML_OP_MUL, "sd.modulate_convrot.vulkan.v1"},
{GGML_OP_GLU, "sd.swiglu_convrot.vulkan.v1"},
{GGML_OP_UNARY, "sd.sigmoid_convrot.vulkan.v1"}};
for (const auto& fusion : fusions) {
d.op = fusion.op;
d.name = fusion.name;
GGML_ASSERT(ggml_ext_register_kernel(device, &d));
}
}
@@ -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();
}
}