Commit af50f2f24 for llama.cpp
commit af50f2f248b475d9760f8e67503d01c0f194f5f4
Author: Aaron Teo <aaron.teo1@ibm.com>
Date: Tue Oct 13 15:53:48 2026 +0800
ggml-zdnn : introduce Q8_0 mat mul (#29836)
* ggml-zdnn: initial prototype of q8_0 matmul
Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>
* docs: update zdnn docs
Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>
---------
Signed-off-by: Aaron Teo <aaron.teo1@ibm.com>
diff --git a/docs/backend/zDNN.md b/docs/backend/zDNN.md
index 94108f5ef..947be0515 100644
--- a/docs/backend/zDNN.md
+++ b/docs/backend/zDNN.md
@@ -27,6 +27,9 @@ The llama.cpp zDNN backend is designed to enable llama.cpp on IBM z17 and later
| F32 | Supported |
| F16 | Supported |
| BF16 | Supported |
+| Q8_0 | Supported |
+
+Note: Quantized weights are requantized to int8 with one scale per group of up to 2048 values in each row (or 1024 with `GGML_ZDNN_HIGH_PREC=1`), after each input channel is divided by its RMS (the activations are multiplied by it instead), and each activation is quantized to two int8 values (about 15 bits). This requires `NNPA_PARMBLKFORMAT_1` (IBM z17 / LinuxONE 5).
## CMake Options
@@ -37,6 +40,12 @@ The IBM zDNN backend has the following CMake options that control the behaviour
| `GGML_ZDNN` | `OFF` | Compile llama.cpp with zDNN support |
| `ZDNN_ROOT` | `""` | Override zDNN library lookup |
+## Environment Variables
+
+| Environment Variable | Default Value | Description |
+| --------------------- | ------------- | -------------------------------------------------------------------------------------- |
+| `GGML_ZDNN_HIGH_PREC` | `0` | Set to `1` to requantize Q8_0 weights in smaller groups, for lower loss at lower speed |
+
## 1. Install zDNN Library
Note: Using the zDNN library provided via `apt` or `yum` may not work correctly as reported in [#15772](https://github.com/ggml-org/llama.cpp/issues/15772). It is preferred that you compile from source.
diff --git a/ggml/src/ggml-zdnn/common.hpp b/ggml/src/ggml-zdnn/common.hpp
index 2462ded55..80a9eadd1 100644
--- a/ggml/src/ggml-zdnn/common.hpp
+++ b/ggml/src/ggml-zdnn/common.hpp
@@ -6,11 +6,17 @@
#include "zdnn.h"
-#include <vector>
+#include <algorithm>
+#include <cstdlib>
+#include <future>
+#include <map>
#include <memory>
+#include <utility>
+#include <vector>
-#define GGML_ZDNN_NAME "zDNN"
-#define GGML_ZDNN_VERSION ZDNN_VERNUM
+#define GGML_ZDNN_NAME "zDNN"
+#define GGML_ZDNN_VERSION ZDNN_VERNUM
+#define GGML_ZDNN_SCRATCH_SLOTS 3 // int8 inputs, DFloat16 partial results, partial results for grouped matmul
#define ZDNN_CHECK(stmt) \
do { \
@@ -30,11 +36,6 @@ struct ggml_backend_zdnn_device_context {
char name[128];
};
-struct ggml_backend_zdnn_context {
- int device;
- ggml_cgraph * gf;
-};
-
struct ggml_backend_zdnn_buffer {
void * data;
ggml_backend_zdnn_buffer * extra; // for bias, etc.
@@ -44,6 +45,9 @@ struct ggml_backend_zdnn_buffer {
zdnn_tensor_desc tfm_desc;
zdnn_ztensor ztensor;
+ std::vector<float> scales; // per-row scales for quantized weights
+ std::vector<float> channel_scales; // per-input-channel scales for quantized weights, applied to the inputs
+
char name[GGML_MAX_NAME];
};
@@ -56,4 +60,75 @@ struct ggml_backend_zdnn_buffer_context {
std::vector<std::unique_ptr<ggml_backend_zdnn_buffer>> buffers;
};
+struct ggml_backend_zdnn_context {
+ int device;
+ int32_t n_threads; // for the CPU work around zDNN calls
+ ggml_cgraph * gf;
+
+ // zero bias for stacked quantized matmuls, keyed by (n_stacks, ne0)
+ std::map<std::pair<int64_t, int64_t>, std::unique_ptr<ggml_backend_zdnn_buffer>> zero_bias;
+
+ std::vector<int8_t> scratch_q;
+ std::vector<float> scratch_f32;
+
+ void * scratch_ztensor[GGML_ZDNN_SCRATCH_SLOTS] = {}; // 4K aligned
+ size_t scratch_ztensor_size[GGML_ZDNN_SCRATCH_SLOTS] = {};
+};
+
+template <typename T>
+static inline T * ggml_zdnn_resize_scratch_max(std::vector<T> & scratch, int64_t n) {
+ if ((int64_t)scratch.size() < n) {
+ scratch.resize(n);
+ }
+
+ return scratch.data();
+}
+
+//
+// Scratch buffers
+//
+
+static inline void ggml_zdnn_init_scratch_ztensor(ggml_backend_zdnn_context * ctx,
+ zdnn_tensor_desc * pre_tfm_desc,
+ zdnn_tensor_desc * tfm_desc,
+ zdnn_ztensor * ztensor,
+ int32_t slot) {
+
+ const uint64_t size = zdnn_getsize_ztensor(tfm_desc);
+
+ if (ctx->scratch_ztensor_size[slot] < size) {
+ free(ctx->scratch_ztensor[slot]);
+ ctx->scratch_ztensor[slot] = std::aligned_alloc(4096, GGML_PAD(size, 4096));
+ ctx->scratch_ztensor_size[slot] = GGML_PAD(size, 4096);
+ GGML_ASSERT(ctx->scratch_ztensor[slot] != nullptr);
+ }
+
+ zdnn_init_quantized_ztensor(pre_tfm_desc, tfm_desc, 1.0f, 0.0f, ztensor);
+ ztensor->buffer = ctx->scratch_ztensor[slot];
+ ztensor->buffer_size = size;
+}
+
+//
+// Threading
+//
+
+// split [0, n) into n_threads ranges, the calling thread runs the first range
+template <typename F>
+static inline void ggml_zdnn_parallel_for(int32_t n_threads, int64_t n, const F & f) {
+ n_threads = (int32_t)std::max<int64_t>(1, std::min<int64_t>(n_threads, n));
+
+ std::vector<std::future<void>> tasks;
+ for (int32_t ith = 1; ith < n_threads; ith++) {
+ const int64_t start = (ith + 0) * n / n_threads;
+ const int64_t end = (ith + 1) * n / n_threads;
+ tasks.push_back(std::async(std::launch::async, [&f, start, end]() { f(start, end); }));
+ }
+
+ f(0, n/n_threads);
+
+ for (auto & task : tasks) {
+ task.get();
+ }
+}
+
#endif // GGML_ZDNN_COMMON_HPP
diff --git a/ggml/src/ggml-zdnn/ggml-zdnn.cpp b/ggml/src/ggml-zdnn/ggml-zdnn.cpp
index 5d1209357..3b717a8c9 100644
--- a/ggml/src/ggml-zdnn/ggml-zdnn.cpp
+++ b/ggml/src/ggml-zdnn/ggml-zdnn.cpp
@@ -4,6 +4,7 @@
#include "ggml-zdnn/common.hpp"
#include "ggml-zdnn/mmf.hpp"
+#include "ggml-zdnn/mmq.hpp"
#include "ggml-zdnn/utils.hpp"
#include "ggml.h"
@@ -13,15 +14,17 @@
#include <unistd.h>
static void ggml_zdnn_compute_forward_mul_mat(
- const ggml_backend_zdnn_context * ctx,
+ ggml_backend_zdnn_context * ctx,
ggml_tensor * dst) {
const ggml_tensor * src0 = dst->src[0]; // weights
const ggml_tensor * src1 = dst->src[1]; // inputs
- // TODO: implement support for quantized types
- // we currently only support f32, f16, and bf16
- ggml_zdnn_mul_mat_f(ctx, src0, src1, dst);
+ if (ggml_is_quantized(src0->type)) {
+ ggml_zdnn_mul_mat_q(ctx, src0, src1, dst);
+ } else {
+ ggml_zdnn_mul_mat_f(ctx, src0, src1, dst);
+ }
}
static bool ggml_zdnn_compute_forward(
@@ -98,7 +101,7 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
if (!ggml_is_matrix(weights) || !ggml_is_matrix(inputs) ||
!ggml_is_contiguous(weights) || !ggml_is_contiguous(inputs) ||
- weights->view_src != nullptr || inputs->view_src != nullptr ||
+ weights->view_src != nullptr ||
ne0 > max_batch || ne1 > max_batch || ne10 > max_batch) {
return false;
}
@@ -107,7 +110,11 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
case GGML_TYPE_F32:
case GGML_TYPE_F16:
case GGML_TYPE_BF16:
- return true;
+ // the inputs ztensor is used, and views have none
+ return inputs->view_src == nullptr;
+ case GGML_TYPE_Q8_0:
+ // the inputs are quantized from their data, so views work as well
+ return ctx_dev->has_parmblkformat_1 && inputs->type == GGML_TYPE_F32;
default:
return false;
}
@@ -186,6 +193,7 @@ static ggml_backend_zdnn_context * ggml_zdnn_init(ggml_backend_dev_t dev) {
GGML_LOG_INFO("%s: NNPA_PARMBLKFORMAT_0 = %s\n", __func__, ctx_dev->has_parmblkformat_0 ? "true" : "false");
GGML_LOG_INFO("%s: NNPA_PARMBLKFORMAT_1 = %s\n", __func__, ctx_dev->has_parmblkformat_1 ? "true" : "false");
+ ctx->n_threads = GGML_DEFAULT_N_THREADS;
ctx->gf = nullptr;
return ctx;
@@ -193,6 +201,14 @@ static ggml_backend_zdnn_context * ggml_zdnn_init(ggml_backend_dev_t dev) {
static void ggml_zdnn_free(ggml_backend_zdnn_context * ctx) {
GGML_LOG_INFO("%s: deallocating\n", __func__);
+
+ for (const auto & it : ctx->zero_bias) {
+ ZDNN_CHECK(zdnn_free_ztensor_buffer(&it.second->ztensor));
+ }
+ for (void * scratch : ctx->scratch_ztensor) {
+ free(scratch);
+ }
+
delete ctx;
}
@@ -295,9 +311,19 @@ static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, gg
ggml_backend_zdnn_buffer * extra = (ggml_backend_zdnn_buffer *)tensor->extra;
+ if (ggml_is_quantized(tensor->type)) {
+ GGML_ASSERT(offset == 0 && size == ggml_nbytes(tensor));
+ ggml_zdnn_load_tensor_q(extra, tensor);
+ return;
+ }
+
// Fixes the LLAMA_SET_ROWS bug
// see: https://github.com/ggml-org/llama.cpp/issues/15414
- if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE && extra->ztensor.is_transformed) zdnn_reset_ztensor(&extra->ztensor);
+ // compute tensors are transformed on first use instead, as not every op reads the ztensor
+ if (tensor->buffer->usage == GGML_BACKEND_BUFFER_USAGE_COMPUTE) {
+ zdnn_reset_ztensor(&extra->ztensor);
+ return;
+ }
if (extra->ztensor.is_transformed == false) ggml_zdnn_load_tensor(extra->ztensor, tensor->data);
GGML_UNUSED(buffer);
@@ -621,10 +647,20 @@ static ggml_backend_feature * ggml_backend_zdnn_get_features(ggml_backend_reg_t
GGML_UNUSED(reg);
}
+static void ggml_backend_zdnn_set_n_threads(ggml_backend_t backend, int n_threads) {
+ GGML_ASSERT(ggml_backend_is_zdnn(backend));
+
+ ggml_backend_zdnn_context * ctx = (ggml_backend_zdnn_context *)backend->context;
+ ctx->n_threads = n_threads;
+}
+
static void * ggml_backend_zdnn_get_proc_address(ggml_backend_reg_t reg, const char * name) {
if (strcmp(name, "ggml_backend_get_features") == 0) {
return (void *) ggml_backend_zdnn_get_features;
}
+ if (strcmp(name, "ggml_backend_set_n_threads") == 0) {
+ return (void *) ggml_backend_zdnn_set_n_threads;
+ }
return NULL;
diff --git a/ggml/src/ggml-zdnn/mmf.cpp b/ggml/src/ggml-zdnn/mmf.cpp
index 3ac9cf3c9..0551fb88b 100644
--- a/ggml/src/ggml-zdnn/mmf.cpp
+++ b/ggml/src/ggml-zdnn/mmf.cpp
@@ -1,5 +1,6 @@
#include "ggml.h"
#include "mmf.hpp"
+#include "utils.hpp"
void ggml_zdnn_mul_mat_f(
const ggml_backend_zdnn_context * ctx,
@@ -65,6 +66,14 @@ void ggml_zdnn_mul_mat_f(
GGML_ASSERT(inputs_extra->pre_tfm_desc.dim1 == inputs->ne[0] && "inputs_extra->pre_tfm_desc.dim1 must match inputs->ne[0]");
GGML_ASSERT(inputs_extra->pre_tfm_desc.dim2 == inputs->ne[1] && "inputs_extra->pre_tfm_desc.dim2 must match inputs->ne[1]");
+ // compute tensors are transformed on first use
+ if (!weights_extra->ztensor.is_transformed) {
+ ggml_zdnn_load_tensor(weights_extra->ztensor, weights->data);
+ }
+ if (!inputs_extra->ztensor.is_transformed) {
+ ggml_zdnn_load_tensor(inputs_extra->ztensor, inputs->data);
+ }
+
ZDNN_CHECK(zdnn_matmul_transpose_op(&inputs_extra->ztensor, &weights_extra->ztensor, &bias_extra->ztensor,
false, true, MATMUL_OP_ADDITION, &output_extra->ztensor));
// TODO: Remove in the future as we are currently DLF16 -> FP32 then in the next op, FP32 -> DLF16 again. Inefficient.
diff --git a/ggml/src/ggml-zdnn/mmq.cpp b/ggml/src/ggml-zdnn/mmq.cpp
new file mode 100644
index 000000000..947ed44c8
--- /dev/null
+++ b/ggml/src/ggml-zdnn/mmq.cpp
@@ -0,0 +1,331 @@
+#include "ggml.h"
+#include "mmq.hpp"
+#include "quantize.hpp"
+
+#include <algorithm>
+#include <future>
+#include <vector>
+
+// view of the stacks [s0, s0 + n) of a stacked ztensor, the stacks are the outermost dim of its buffer
+// the view descs must outlive the view
+static zdnn_ztensor ggml_zdnn_stack_view(zdnn_tensor_desc * pre_tfm_desc,
+ zdnn_tensor_desc * tfm_desc,
+ const zdnn_ztensor & src,
+ int64_t n_stacks,
+ int64_t s0,
+ int64_t n) {
+
+ GGML_ASSERT(src.buffer_size % n_stacks == 0);
+ const uint64_t stack_bytes = src.buffer_size / n_stacks;
+
+ *pre_tfm_desc = *src.pre_transformed_desc;
+ *tfm_desc = *src.transformed_desc;
+
+ switch (pre_tfm_desc->layout) {
+ case ZDNN_3DS:
+ {
+ pre_tfm_desc->dim3 = n;
+ } break;
+ case ZDNN_2DS:
+ {
+ pre_tfm_desc->dim2 = n;
+ } break;
+ default:
+ GGML_ABORT("%s: unsupported stacked layout", __func__);
+ }
+
+ tfm_desc->dim4 = n;
+ GGML_ASSERT(zdnn_getsize_ztensor(tfm_desc) == n*stack_bytes);
+
+ zdnn_ztensor view = src;
+ view.pre_transformed_desc = pre_tfm_desc;
+ view.transformed_desc = tfm_desc;
+ view.buffer = (char *) src.buffer + s0*stack_bytes;
+ view.buffer_size = n * stack_bytes;
+ return view;
+}
+
+// zero bias of n_stacks stacks, created once per shape as it never changes
+static const zdnn_ztensor & ggml_zdnn_zero_bias(ggml_backend_zdnn_context * ctx, int64_t n_stacks, int64_t ne0) {
+ std::unique_ptr<ggml_backend_zdnn_buffer> & bias = ctx->zero_bias[{ n_stacks, ne0 }];
+
+ if (!bias) {
+ bias = std::make_unique<ggml_backend_zdnn_buffer>();
+
+ std::vector<float> zeros(n_stacks * ne0, 0.0f);
+
+ zdnn_init_pre_transformed_desc(ZDNN_2DS, FP32, &bias->pre_tfm_desc, n_stacks, ne0);
+ ZDNN_CHECK(zdnn_generate_transformed_desc(&bias->pre_tfm_desc, &bias->tfm_desc));
+ ZDNN_CHECK(zdnn_init_ztensor_with_malloc(&bias->pre_tfm_desc, &bias->tfm_desc, &bias->ztensor));
+ ZDNN_CHECK(zdnn_transform_ztensor(&bias->ztensor, zeros.data()));
+ }
+
+ return bias->ztensor;
+}
+
+// the weights are stacked into groups along ne00, and each group of each row of the inputs and weights has its own scale
+// zDNN multiplies the groups in stacked calls, then the partial results of each group are scaled and summed here
+static void ggml_zdnn_mul_mat_q_grouped(ggml_backend_zdnn_context * ctx,
+ int32_t n_threads,
+ const ggml_tensor * inputs,
+ const ggml_backend_zdnn_buffer * weights_extra,
+ ggml_tensor * output) {
+ const int64_t ne10 = inputs->ne[0];
+ const int64_t ne11 = inputs->ne[1];
+ const int64_t ne0 = output->ne[0];
+
+ const int64_t n_groups = weights_extra->pre_tfm_desc.dim3;
+ const int64_t gs = weights_extra->pre_tfm_desc.dim2;
+
+ // the last group is zero padded when gs does not divide ne10
+ GGML_ASSERT(n_groups * gs >= ne10 && (n_groups - 1) * gs < ne10);
+
+ // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
+ int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * n_groups * gs * ne11);
+ std::vector<float> inputs_scales(n_groups * ne11);
+ ggml_zdnn_quantize_inputs(n_threads, inputs, gs, weights_extra->channel_scales.data(),
+ inputs_q, inputs_q + n_groups * gs * ne11, inputs_scales.data());
+
+ zdnn_tensor_desc inputs_pre_tfm_desc;
+ zdnn_tensor_desc inputs_tfm_desc;
+ zdnn_ztensor inputs_ztensor;
+
+ zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2 * n_groups, ne11, gs);
+ ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
+ ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+ ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
+
+ // each group of the partial results is read back as its own 2D ztensor
+ zdnn_tensor_desc group_pre_tfm_desc;
+ zdnn_tensor_desc group_tfm_desc;
+
+ zdnn_init_pre_transformed_desc(ZDNN_2D, FP32, &group_pre_tfm_desc, ne11, ne0);
+ ZDNN_CHECK(zdnn_generate_transformed_desc(&group_pre_tfm_desc, &group_tfm_desc));
+
+ const uint64_t group_bytes = zdnn_getsize_ztensor(&group_tfm_desc);
+
+ // stacked calls with a large output are much slower per group, so run the groups in chunks
+ // 8 MiB per chunk was benchmarked to be the best on Llama-3.2-1B Q8_0
+ // always re-measure with different models to see if different sizes need a dispatching logic
+ constexpr int64_t chunk_bytes = 8*1024*1024;
+
+ const int64_t n_chunk = std::clamp<int64_t>(chunk_bytes / (int64_t)group_bytes, 1, n_groups);
+
+ // all zeros, so it is already the pre-computed bias for every chunk
+ const zdnn_ztensor & bias_ztensor = ggml_zdnn_zero_bias(ctx, n_chunk, ne0);
+
+ // two partial buffers, so the next chunk runs on the NNPA while this one is summed
+ // each holds the results of a chunk of inputs, followed by the results of their rounding errors
+ zdnn_tensor_desc partial_pre_tfm_desc;
+ zdnn_tensor_desc partial_tfm_desc;
+ zdnn_ztensor partial_ztensor[2];
+
+ zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2 * n_chunk, ne11, ne0);
+ ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
+ GGML_ASSERT(zdnn_getsize_ztensor(&partial_tfm_desc) == (uint64_t)(2 * n_chunk) * group_bytes);
+ for (int32_t i = 0; i < 2; i++) {
+ ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor[i], 1 + i);
+ }
+
+ const float * weights_scales = weights_extra->scales.data();
+ float * partial = ggml_zdnn_resize_scratch_max(ctx->scratch_f32, 2 * ne11 * ne0);
+ float * partial_res = partial + ne11 * ne0;
+
+ const auto sum_chunk = [&](const zdnn_ztensor & chunk, int64_t g0, int64_t n) {
+ zdnn_ztensor group_ztensor;
+ zdnn_init_ztensor(&group_pre_tfm_desc, &group_tfm_desc, &group_ztensor);
+
+ for (int64_t g = g0; g < g0 + n; g++) {
+ group_ztensor.buffer_size = group_bytes;
+ group_ztensor.is_transformed = true;
+
+ group_ztensor.buffer = (char *)chunk.buffer + (g - g0)*group_bytes;
+ ZDNN_CHECK(zdnn_transform_origtensor(&group_ztensor, partial));
+
+ group_ztensor.buffer = (char *)chunk.buffer + (n_chunk + g - g0)*group_bytes;
+ ZDNN_CHECK(zdnn_transform_origtensor(&group_ztensor, partial_res));
+
+ const float * sw = weights_scales + g*ne0;
+ for (int64_t i1 = 0; i1 < ne11; i1++) {
+ const float sa = inputs_scales[g*ne11 + i1];
+ const float * p = partial + i1*ne0;
+ const float * pr = partial_res + i1*ne0;
+ float * y = (float *)output->data + i1*ne0;
+
+ if (g == 0) {
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ y[i0] = sa * sw[i0] * (p[i0] + pr[i0] * (1.0f/254.0f));
+ }
+ } else {
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ y[i0] += sa * sw[i0] * (p[i0] + pr[i0] * (1.0f/254.0f));
+ }
+ }
+ }
+ }
+ };
+
+ zdnn_tensor_desc view_pre_tfm_desc[6];
+ zdnn_tensor_desc view_tfm_desc[6];
+ std::future<void> sum_task;
+
+ for (int64_t c = 0, g0 = 0; g0 < n_groups; c++, g0 += n_chunk) {
+ const int64_t n = std::min(n_chunk, n_groups - g0);
+
+ const zdnn_ztensor inputs_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[0], &view_tfm_desc[0], inputs_ztensor, 2 * n_groups, g0, n);
+ const zdnn_ztensor inputs_res_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[1], &view_tfm_desc[1], inputs_ztensor, 2 * n_groups, n_groups + g0, n);
+ const zdnn_ztensor weights_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[2], &view_tfm_desc[2], weights_extra->ztensor, n_groups, g0, n);
+ const zdnn_ztensor bias_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[3], &view_tfm_desc[3], bias_ztensor, n_chunk, 0, n);
+
+ // the partial buffer of chunk c - 2 is summed by now, as sum_task is waited on before starting the next sum
+ zdnn_ztensor & chunk = partial_ztensor[c % 2];
+ chunk.pre_transformed_desc = &partial_pre_tfm_desc;
+ chunk.transformed_desc = &partial_tfm_desc;
+ chunk.is_transformed = false;
+
+ zdnn_ztensor chunk_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[4], &view_tfm_desc[4], chunk, 2 * n_chunk, 0, n);
+ zdnn_ztensor chunk_res_view = ggml_zdnn_stack_view(&view_pre_tfm_desc[5], &view_tfm_desc[5], chunk, 2 * n_chunk, n_chunk, n);
+ ZDNN_CHECK(zdnn_quantized_matmul_op(&inputs_view, &weights_view, &bias_view,
+ MATMUL_OP_ADDITION, -127, 127, true, false, true,
+ nullptr, &chunk_view));
+ ZDNN_CHECK(zdnn_quantized_matmul_op(&inputs_res_view, &weights_view, &bias_view,
+ MATMUL_OP_ADDITION, -127, 127, true, false, true,
+ nullptr, &chunk_res_view));
+
+ if (sum_task.valid()) {
+ sum_task.get();
+ }
+
+ if (n_threads > 1) {
+ sum_task = std::async(std::launch::async, sum_chunk, chunk, g0, n);
+ } else {
+ sum_chunk(chunk, g0, n);
+ }
+ }
+
+ if (sum_task.valid()) {
+ sum_task.get();
+ }
+
+}
+
+void ggml_zdnn_mul_mat_q(
+ ggml_backend_zdnn_context * ctx,
+ const ggml_tensor * src0,
+ const ggml_tensor * src1,
+ ggml_tensor * dst) {
+ GGML_TENSOR_BINARY_OP_LOCALS;
+
+ const enum ggml_type type = src0->type;
+
+ GGML_ASSERT(ne0 == ne01);
+ GGML_ASSERT(ne1 == ne11);
+ GGML_ASSERT(ne2 == ne12);
+ GGML_ASSERT(ne3 == ne13);
+
+ // we don't support permuted src0 or src1
+ GGML_ASSERT(nb00 == ggml_type_size(type));
+ GGML_ASSERT(nb10 == ggml_type_size(src1->type));
+
+ // dst cannot be transposed or permuted
+ GGML_ASSERT(nb0 == sizeof(float));
+ GGML_ASSERT(nb0 <= nb1);
+ GGML_ASSERT(nb1 <= nb2);
+ GGML_ASSERT(nb2 <= nb3);
+
+ const ggml_tensor * weights = src0;
+ const ggml_tensor * inputs = src1;
+ ggml_tensor * output = dst;
+
+ GGML_ASSERT(inputs->type == GGML_TYPE_F32);
+
+ ggml_backend_zdnn_buffer * weights_extra = (ggml_backend_zdnn_buffer *)weights->extra;
+ ggml_backend_zdnn_buffer * output_extra = (ggml_backend_zdnn_buffer *)output->extra;
+
+ // starting threads costs more than it saves for small inputs, e.g. during token generation
+ const int32_t n_threads = ne10*ne11 >= 65536 ? ctx->n_threads : 1;
+
+ if (weights_extra->pre_tfm_desc.layout == ZDNN_3DS) {
+ GGML_ASSERT(weights_extra->pre_tfm_desc.dim1 == weights->ne[1] && "weights_extra->pre_tfm_desc.dim1 must match weights->ne[1]");
+
+ ggml_zdnn_mul_mat_q_grouped(ctx, n_threads, inputs, weights_extra, output);
+
+ // the result is written to output->data directly, so the output ztensor is stale
+ output_extra->ztensor.is_transformed = false;
+ return;
+ }
+
+ ggml_backend_zdnn_buffer * bias_extra = (ggml_backend_zdnn_buffer *)output_extra->extra;
+
+ GGML_ASSERT(weights_extra->pre_tfm_desc.dim2 == weights->ne[0] && "weights_extra->pre_tfm_desc.dim2 must match weights->ne[0]");
+ GGML_ASSERT(weights_extra->pre_tfm_desc.dim1 == weights->ne[1] && "weights_extra->pre_tfm_desc.dim1 must match weights->ne[1]");
+
+ // zDNN takes one scale per zTensor, so quantize each row of the inputs with its own scale
+ // the rounding errors of the inputs are stacked after the inputs, see ggml_zdnn_quantize_inputs
+ int8_t * inputs_q = ggml_zdnn_resize_scratch_max(ctx->scratch_q, 2 * ne10 * ne11);
+ std::vector<float> inputs_scales(ne11);
+ ggml_zdnn_quantize_inputs(n_threads, inputs, ne10, weights_extra->channel_scales.data(),
+ inputs_q, inputs_q + ne10 * ne11, inputs_scales.data());
+
+ zdnn_tensor_desc inputs_pre_tfm_desc;
+ zdnn_tensor_desc inputs_tfm_desc;
+ zdnn_ztensor inputs_ztensor;
+
+ zdnn_init_pre_transformed_desc(ZDNN_3DS, INT8, &inputs_pre_tfm_desc, 2, ne11, ne10);
+ ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&inputs_pre_tfm_desc, QUANTIZED_INT8, &inputs_tfm_desc));
+ ggml_zdnn_init_scratch_ztensor(ctx, &inputs_pre_tfm_desc, &inputs_tfm_desc, &inputs_ztensor, 0);
+ ZDNN_CHECK(zdnn_transform_quantized_ztensor(&inputs_ztensor, false, -127, 127, inputs_q));
+
+ zdnn_tensor_desc partial_pre_tfm_desc;
+ zdnn_tensor_desc partial_tfm_desc;
+ zdnn_ztensor partial_ztensor;
+
+ zdnn_init_pre_transformed_desc(ZDNN_3DS, FP32, &partial_pre_tfm_desc, 2, ne11, ne0);
+ ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&partial_pre_tfm_desc, QUANTIZED_DLFLOAT16, &partial_tfm_desc));
+ ggml_zdnn_init_scratch_ztensor(ctx, &partial_pre_tfm_desc, &partial_tfm_desc, &partial_ztensor, 1);
+
+ // both stacks are multiplied by the same unstacked weights in one broadcast call
+ // the bias is all zeros, so it is already the pre-computed bias
+ ZDNN_CHECK(zdnn_quantized_matmul_op(&inputs_ztensor, &weights_extra->ztensor, &bias_extra->ztensor,
+ MATMUL_OP_ADDITION, -127, 127, true, false, true,
+ nullptr, &partial_ztensor));
+
+ // each stack of the partial results is read back as its own 2D ztensor
+ zdnn_tensor_desc stack_pre_tfm_desc;
+ zdnn_tensor_desc stack_tfm_desc;
+
+ zdnn_init_pre_transformed_desc(ZDNN_2D, FP32, &stack_pre_tfm_desc, ne11, ne0);
+ ZDNN_CHECK(zdnn_generate_transformed_desc(&stack_pre_tfm_desc, &stack_tfm_desc));
+
+ const uint64_t stack_bytes = zdnn_getsize_ztensor(&stack_tfm_desc);
+ GGML_ASSERT(partial_ztensor.buffer_size == 2*stack_bytes);
+
+ float * partial[2] = { (float *)output->data, ggml_zdnn_resize_scratch_max(ctx->scratch_f32, ne11 * ne0) };
+
+ ggml_zdnn_parallel_for(n_threads, 2, [&](int64_t s_start, int64_t s_end) {
+ zdnn_ztensor stack_ztensor;
+ zdnn_init_ztensor(&stack_pre_tfm_desc, &stack_tfm_desc, &stack_ztensor);
+
+ for (int64_t s = s_start; s < s_end; s++) {
+ stack_ztensor.buffer = (char *)partial_ztensor.buffer + s*stack_bytes;
+ stack_ztensor.buffer_size = stack_bytes;
+ stack_ztensor.is_transformed = true;
+ ZDNN_CHECK(zdnn_transform_origtensor(&stack_ztensor, partial[s]));
+ }
+ });
+
+ // the result is written to output->data directly, so the output ztensor is stale
+ output_extra->ztensor.is_transformed = false;
+
+ const float * weights_scales = weights_extra->scales.data();
+ ggml_zdnn_parallel_for(n_threads, ne1, [&](int64_t i1_start, int64_t i1_end) {
+ for (int64_t i1 = i1_start; i1 < i1_end; i1++) {
+ const float sa = inputs_scales[i1];
+ const float * pr = partial[1] + i1*ne0;
+ float * y = partial[0] + i1*ne0;
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ y[i0] = sa * weights_scales[i0] * (y[i0] + pr[i0] * (1.0f/254.0f));
+ }
+ }
+ });
+}
diff --git a/ggml/src/ggml-zdnn/mmq.hpp b/ggml/src/ggml-zdnn/mmq.hpp
new file mode 100644
index 000000000..3e49f0af8
--- /dev/null
+++ b/ggml/src/ggml-zdnn/mmq.hpp
@@ -0,0 +1,12 @@
+#ifndef GGML_ZDNN_MMQ_HPP
+#define GGML_ZDNN_MMQ_HPP
+
+#include "common.hpp"
+
+void ggml_zdnn_mul_mat_q(
+ ggml_backend_zdnn_context * ctx,
+ const ggml_tensor * src0,
+ const ggml_tensor * src1,
+ ggml_tensor * dst);
+
+#endif // GGML_ZDNN_MMQ_HPP
diff --git a/ggml/src/ggml-zdnn/quantize.cpp b/ggml/src/ggml-zdnn/quantize.cpp
new file mode 100644
index 000000000..5012e92e9
--- /dev/null
+++ b/ggml/src/ggml-zdnn/quantize.cpp
@@ -0,0 +1,58 @@
+#include "ggml.h"
+#include "quantize.hpp"
+
+#include <algorithm>
+#include <cmath>
+
+// quantize each group of group_size values per row to int8, plus the rounding error into residuals for ~15 bits
+void ggml_zdnn_quantize_inputs(int32_t n_threads,
+ const ggml_tensor * inputs,
+ int64_t group_size,
+ const float * channel_scales,
+ int8_t * quants,
+ int8_t * residuals,
+ float * scales) {
+
+ const int64_t ne0 = inputs->ne[0];
+ const int64_t ne1 = inputs->ne[1];
+ const int64_t n_groups = (ne0 + group_size - 1) / group_size;
+
+ ggml_zdnn_parallel_for(n_threads, ne1, [&](int64_t i1_start, int64_t i1_end) {
+ for (int64_t i1 = i1_start; i1 < i1_end; i1++) {
+ const float * xi = (const float *)((const char *)inputs->data + i1*inputs->nb[1]);
+
+ for (int64_t g = 0; g < n_groups; g++) {
+ const float * xg = xi + g*group_size;
+ const float * cg = channel_scales + g*group_size;
+ const int64_t n = std::min(group_size, ne0 - g*group_size);
+
+ float amax = 0.0f;
+ for (int64_t i0 = 0; i0 < n; i0++) {
+ amax = std::max(amax, fabsf(xg[i0] * cg[i0]));
+ }
+
+ // the rounding error is at most d/2
+ const float d = amax / 127.0f;
+ const float id = d ? 1.0f/d : 0.0f;
+ const float id_res = 254.0f * id;
+ scales[g*ne1 + i1] = d;
+
+ int8_t * qg = quants + (g*ne1 + i1)*group_size;
+ int8_t * rg = residuals + (g*ne1 + i1)*group_size;
+
+ for (int64_t i0 = 0; i0 < n; i0++) {
+ const float xc = xg[i0] * cg[i0];
+ const float v = roundf(xc * id);
+
+ qg[i0] = (int8_t)v;
+ rg[i0] = (int8_t)std::clamp(roundf((xc - v * d) * id_res), -127.0f, 127.0f);
+ }
+
+ for (int64_t i0 = n; i0 < group_size; i0++) {
+ qg[i0] = 0;
+ rg[i0] = 0;
+ }
+ }
+ }
+ });
+}
diff --git a/ggml/src/ggml-zdnn/quantize.hpp b/ggml/src/ggml-zdnn/quantize.hpp
new file mode 100644
index 000000000..51acbe8ad
--- /dev/null
+++ b/ggml/src/ggml-zdnn/quantize.hpp
@@ -0,0 +1,14 @@
+#ifndef GGML_ZDNN_QUANTIZE_HPP
+#define GGML_ZDNN_QUANTIZE_HPP
+
+#include "common.hpp"
+
+void ggml_zdnn_quantize_inputs(int32_t n_threads,
+ const ggml_tensor * inputs,
+ int64_t group_size,
+ const float * channel_scales,
+ int8_t * quants,
+ int8_t * residuals,
+ float * scales);
+
+#endif // GGML_ZDNN_QUANTIZE_HPP
diff --git a/ggml/src/ggml-zdnn/utils.cpp b/ggml/src/ggml-zdnn/utils.cpp
index 2977cb0fe..18e9879d0 100644
--- a/ggml/src/ggml-zdnn/utils.cpp
+++ b/ggml/src/ggml-zdnn/utils.cpp
@@ -1,6 +1,10 @@
#include "ggml.h"
#include "utils.hpp"
+#include <algorithm>
+#include <cmath>
+#include <cstdlib>
+
zdnn_data_types ggml_zdnn_type_mapping(ggml_type type) {
switch (type) {
case GGML_TYPE_F32:
@@ -9,8 +13,6 @@ zdnn_data_types ggml_zdnn_type_mapping(ggml_type type) {
return FP16;
case GGML_TYPE_BF16:
return BFLOAT;
- case GGML_TYPE_Q8_0:
- return INT8;
case GGML_TYPE_I8:
return INT8;
case GGML_TYPE_I32:
@@ -43,7 +45,141 @@ void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer) {
ZDNN_CHECK(zdnn_transform_ztensor(&ztensor, buffer));
}
+// size of the groups along ne0 that get their own scale, or 0 for one scale per row
+// the groups are spread evenly over ne0 and rounded up to 128 values, the last group is zero padded
+int64_t ggml_zdnn_q_group_size(const ggml_tensor * tensor) {
+ // benchmarked on Llama-3.2-1B and Qwen2.5-1.5B Q8_0 with the channel scales:
+ // 2048 is 15-18% faster than 1024, and 1024 lowers the KL divergence to the CPU by 13-19%
+ // GGML_ZDNN_HIGH_PREC=1 selects 1024 for users that want the lower loss
+ // always re-measure with different models to see if different sizes need a dispatching logic
+ static const bool high_prec = std::getenv("GGML_ZDNN_HIGH_PREC") != nullptr &&
+ std::atoi(std::getenv("GGML_ZDNN_HIGH_PREC")) != 0;
+ const int64_t group_size_max = high_prec ? 1024 : 2048;
+
+ const int64_t ne0 = tensor->ne[0];
+ if (ne0 <= group_size_max) {
+ return 0;
+ }
+
+ const int64_t n_groups = (ne0 + group_size_max - 1) / group_size_max;
+ return GGML_PAD((ne0 + n_groups - 1) / n_groups, 128);
+}
+
+void ggml_zdnn_load_tensor_q(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor) {
+ const int64_t ne0 = tensor->ne[0];
+ const int64_t ne1 = tensor->ne[1];
+
+ const ggml_to_float_t to_float = ggml_get_type_traits(tensor->type)->to_float;
+ GGML_ASSERT(to_float != nullptr);
+
+ // zDNN takes one scale per zTensor, so requantize each group of each row to int8 with its own scale
+ // a group is the whole row, unless the weights are stacked into groups along ne0
+ // rows are written transposed as zDNN computes inputs (ne11, ne0) x weights (ne0, ne1)
+ const bool stacked = buffer->pre_tfm_desc.layout == ZDNN_3DS;
+ const int64_t n_groups = stacked ? buffer->pre_tfm_desc.dim3 : 1;
+ const int64_t gs = stacked ? buffer->pre_tfm_desc.dim2 : ne0;
+
+ // benchmarked: 256 rows fill a 256-byte cache line per write, padding each row by a line avoids cache set conflicts
+ constexpr int64_t n_tile = 256;
+ const int64_t x_stride = ne0 + 64;
+
+ // the rows past ne0 of the last group stay zero as padding
+ std::vector<int8_t> q(n_groups * gs * ne1);
+ std::vector<float> x(n_tile * x_stride);
+ std::vector<float> id(n_tile * n_groups);
+
+ // a few input channels with larger weights than the rest would set the scale of every group,
+ // so divide each input channel by its RMS over the rows and multiply the inputs by it instead
+ // with groups of 1024, this lowered the KL divergence to the CPU by 10-18% on Llama-3.2-1B and Qwen2.5-1.5B Q8_0 at the same speed
+ std::vector<double> sum_sq(ne0, 0.0);
+ for (int64_t i1 = 0; i1 < ne1; i1++) {
+ to_float((const char *)tensor->data + i1*tensor->nb[1], x.data(), ne0);
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ sum_sq[i0] += (double)x[i0] * x[i0];
+ }
+ }
+
+ std::vector<float> channel_id(ne0);
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ const float c = (float)std::sqrt(sum_sq[i0] / ne1);
+ buffer->channel_scales[i0] = c ? c : 1.0f;
+ channel_id[i0] = 1.0f / buffer->channel_scales[i0];
+ }
+
+ for (int64_t i1 = 0; i1 < ne1; i1 += n_tile) {
+ const int64_t nt = std::min(n_tile, ne1 - i1);
+
+ for (int64_t t = 0; t < nt; t++) {
+ float * xt = x.data() + t*x_stride;
+ to_float((const char *)tensor->data + (i1 + t)*tensor->nb[1], xt, ne0);
+
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ xt[i0] *= channel_id[i0];
+ }
+
+ for (int64_t g = 0; g < n_groups; g++) {
+ float amax = 0.0f;
+ for (int64_t i0 = g*gs; i0 < std::min((g + 1)*gs, ne0); i0++) {
+ amax = std::max(amax, fabsf(xt[i0]));
+ }
+
+ const float d = amax / 127.0f;
+ buffer->scales[g*ne1 + i1 + t] = d;
+ id[t*n_groups + g] = d ? 1.0f/d : 0.0f;
+ }
+ }
+
+ for (int64_t i0 = 0; i0 < ne0; i0++) {
+ const int64_t g = i0 / gs;
+ int8_t * qi = q.data() + i0*ne1 + i1;
+ for (int64_t t = 0; t < nt; t++) {
+ qi[t] = (int8_t)roundf(x[t*x_stride + i0] * id[t*n_groups + g]);
+ }
+ }
+ }
+
+ if (buffer->ztensor.is_transformed) {
+ zdnn_reset_ztensor(&buffer->ztensor);
+ }
+ ZDNN_CHECK(zdnn_transform_quantized_ztensor(&buffer->ztensor, false, -127, 127, q.data()));
+}
+
void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor) {
+ if (ggml_is_quantized(tensor->type)) {
+ GGML_ASSERT(ggml_is_matrix(tensor));
+
+ const int64_t gs = ggml_zdnn_q_group_size(tensor);
+
+ if (gs > 0) {
+ // stack the weights into groups along ne0, each group of each row gets its own scale
+ const int64_t n_groups = (tensor->ne[0] + gs - 1) / gs;
+
+ zdnn_init_pre_transformed_desc(
+ ZDNN_3DS,
+ INT8,
+ &buffer->pre_tfm_desc,
+ n_groups, gs, tensor->ne[1]
+ );
+
+ buffer->scales.resize(n_groups * tensor->ne[1]);
+ } else {
+ zdnn_init_pre_transformed_desc(
+ ZDNN_2D,
+ INT8,
+ &buffer->pre_tfm_desc,
+ tensor->ne[0], tensor->ne[1]
+ );
+
+ buffer->scales.resize(tensor->ne[1]);
+ }
+
+ buffer->channel_scales.resize(tensor->ne[0]);
+
+ ZDNN_CHECK(zdnn_generate_quantized_transformed_desc(&buffer->pre_tfm_desc, QUANTIZED_WEIGHTS_INT8, &buffer->tfm_desc));
+ ZDNN_CHECK(zdnn_init_quantized_ztensor_with_malloc(&buffer->pre_tfm_desc, &buffer->tfm_desc, 1.0f, 0.0f, &buffer->ztensor));
+ return;
+ }
+
switch (tensor->op) {
case GGML_OP_MUL_MAT:
{
diff --git a/ggml/src/ggml-zdnn/utils.hpp b/ggml/src/ggml-zdnn/utils.hpp
index c1e2028ed..20f5be5b4 100644
--- a/ggml/src/ggml-zdnn/utils.hpp
+++ b/ggml/src/ggml-zdnn/utils.hpp
@@ -14,6 +14,10 @@ void ggml_zdnn_create_tensor(zdnn_tensor_desc & pre_tfm_desc,
void ggml_zdnn_load_tensor(zdnn_ztensor & ztensor, void * buffer);
+int64_t ggml_zdnn_q_group_size(const ggml_tensor * tensor);
+
+void ggml_zdnn_load_tensor_q(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor);
+
void ggml_zdnn_init_tensor(ggml_backend_zdnn_buffer * buffer, const ggml_tensor * tensor);
#endif // GGML_ZDNN_UTILITIES_HPP