diff --git a/src/llama.cpp b/src/llama.cpp index bfface76..1ce8c431 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -13,6 +13,8 @@ #include "profiler.h" #include "network-utils.h" +#include "quantization.h" + #ifdef GGML_USE_RPC # include "ggml-rpc.h" #endif @@ -103,78 +105,6 @@ #include #include -#define Q8_0_BLOCK_SIZE 64 - -// One Q8_0 block = scale + 64 int8s -typedef struct { - float d; // scale - int8_t qs[Q8_0_BLOCK_SIZE]; // quantized values -} block_q8_0; - -// Quantize a 1D float array into Q8_0 blocks. -// - in: pointer to N floats -// - N: number of elements -// - out_blocks: *out set to malloc'd array of blocks; caller frees -// Returns number of blocks (ceil(N/32)). On error returns 0. -int64_t q8_0_quantize(const float *in, int64_t N, block_q8_0 **out_blocks) { - if (!in || !out_blocks || N == 0) return 0; - - const int64_t nb = (N + Q8_0_BLOCK_SIZE - 1) / Q8_0_BLOCK_SIZE; - block_q8_0 *blk = (block_q8_0 *)malloc(nb * sizeof(block_q8_0)); - if (!blk) return 0; - - for (int64_t b = 0; b < nb; ++b) { - const int64_t start = b * Q8_0_BLOCK_SIZE; - const int64_t rem = (start + Q8_0_BLOCK_SIZE <= N) ? Q8_0_BLOCK_SIZE : (N - start); - - // 1) find max-abs in this block - float amax = 0.0f; - for (int64_t i = 0; i < rem; ++i) { - float v = fabsf(in[start + i]); - if (v > amax) amax = v; - } - - // 2) compute scale - float d = (amax > 0.0f) ? (amax / 127.0f) : 0.0f; - float invd = (d > 0.0f) ? (1.0f / d) : 0.0f; - blk[b].d = d; - - // 3) quantize present elems, zero-pad the rest - for (int64_t i = 0; i < rem; ++i) { - float r = in[start + i] * invd; - // nearest int, then clamp to [-127, 127] - long qi = lrintf(r); - if (qi < -127) qi = -127; - if (qi > 127) qi = 127; - blk[b].qs[i] = (int8_t)qi; - } - for (int64_t i = rem; i < Q8_0_BLOCK_SIZE; ++i) { - blk[b].qs[i] = 0; - } - } - - *out_blocks = blk; - return nb; -} - -// Dequantize Q8_0 blocks back to floats. -// - blocks: pointer to nb blocks -// - N: number of output floats desired (original length) -// - out: pointer to N floats (must be allocated by caller) -void q8_0_dequantize(const block_q8_0 *blocks, int64_t nb, int64_t N, float *out) { - if (!blocks || !out || N == 0) return; - - for (int64_t b = 0; b < nb; ++b) { - const int64_t start = b * Q8_0_BLOCK_SIZE; - const int64_t rem = (start + Q8_0_BLOCK_SIZE <= N) ? Q8_0_BLOCK_SIZE : (N - start); - const float d = blocks[b].d; - - for (int64_t i = 0; i < rem; ++i) { - out[start + i] = d * (float)blocks[b].qs[i]; - } - } -} - int g_llama_send_tensors_counts = 0; int g_llama_recv_tensors_counts = 0; @@ -18209,42 +18139,42 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba g_llama_send_tensors_counts++; try { std::vector send_msgs; - + int64_t num_elements = tensors->sub_gf_out->ne[0] * tensors->sub_gf_out->ne[1]; + int64_t float_element_size = num_elements * sizeof(float); + quantized_array_t * quantized_array = NULL; + if (quantize(ubatch->backend_embd, num_elements, 0 /*q8_0*/, &quantized_array) || !quantized_array) { + LLAMA_LOG_INFO("Failed to allocate space or doing quantization\n"); + return; + } + int64_t buf_size = get_quantized_array_size(quantized_array); + send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); - send_msgs.emplace_back("q8_0", strlen("q8_0")); + send_msgs.emplace_back("quantized", strlen("quantized")); send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne)); - // size_t buf_size = tensors->sub_gf_out->ne[0] * tensors->sub_gf_out->ne[1] * sizeof(float); - // send_msgs.emplace_back(ubatch->backend_embd, buf_size); - - // q8_0 communication quantization - int64_t tensor_elements = tensors->sub_gf_out->ne[0] * tensors->sub_gf_out->ne[1]; - block_q8_0 * quantized_data = NULL; - int64_t n_blocks = q8_0_quantize(ubatch->backend_embd, tensor_elements, &quantized_data); - int64_t buf_size = n_blocks * sizeof(block_q8_0); - send_msgs.emplace_back(&n_blocks, sizeof(n_blocks)); - send_msgs.emplace_back(quantized_data, buf_size); + send_msgs.emplace_back(quantized_array, buf_size); + send_msgs.emplace_back(&buf_size, sizeof(int64_t)); if (tensors->inp_pos) { - send_msgs.emplace_back("inp_pos", strlen("inp_pos")); - send_msgs.emplace_back("int32", strlen("int32")); - send_msgs.emplace_back(tensors->inp_pos->ne, sizeof(tensors->inp_pos->ne[0])); - buf_size = tensors->inp_pos->ne[0] * sizeof(int32_t); int64_t zero = 0; - send_msgs.emplace_back(&zero, sizeof(int64_t)); // extra frame that recv does not account for + buf_size = tensors->inp_pos->ne[0] * sizeof(int32_t); + + send_msgs.emplace_back("inp_pos", strlen("inp_pos")); + send_msgs.emplace_back("normal", strlen("normal")); + send_msgs.emplace_back(tensors->inp_pos->ne, sizeof(tensors->inp_pos->ne[0])); send_msgs.emplace_back(ubatch->pos, buf_size); + send_msgs.emplace_back(&zero, sizeof(int64_t)); } zmq::send_multipart(socket, send_msgs); - free(quantized_data); + free_quantized_array(quantized_array); - // TODO: Fix dump feature to support int8 - // if (dump_folder && strlen(dump_folder) > 0) { - // std::string dump_path = std::string(dump_folder) + "/send_" + std::to_string(g_llama_send_tensors_counts) + ".bin"; - // dump_tensors(dump_path, static_cast(TensorDataType::FLOAT32), - // static_cast(tensors->sub_gf_out->ne[0]), - // static_cast(tensors->sub_gf_out->ne[1]), - // buf_size, ubatch->backend_embd); - // } + if (dump_folder && strlen(dump_folder) > 0) { + std::string dump_path = std::string(dump_folder) + "/send_" + std::to_string(g_llama_send_tensors_counts) + ".bin"; + dump_tensors(dump_path, static_cast(TensorDataType::FLOAT32), + static_cast(tensors->sub_gf_out->ne[0]), + static_cast(tensors->sub_gf_out->ne[1]), + float_element_size, ubatch->backend_embd); + } } catch (const zmq::error_t& e) { LLAMA_LOG_INFO("Failed to send tensor data: %s\n", e.what()); @@ -18263,37 +18193,32 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba std::string key = recv_msgs[i].to_string(); std::string comm_type = recv_msgs[i + 1].to_string(); zmq::message_t &dims_msg = recv_msgs[i + 2]; - zmq::message_t &n_blocks_msg = recv_msgs[i + 3]; - zmq::message_t &data_msg = recv_msgs[i + 4]; + zmq::message_t &data_msg = recv_msgs[i + 3]; + zmq::message_t &buffer_size_msg = recv_msgs[i + 4]; if (key == "sub_gf_out") { int64_t * dims = static_cast(dims_msg.data()); - int64_t * n_blocks_ptr = static_cast(n_blocks_msg.data()); - int64_t n_blocks = *n_blocks_ptr; - int64_t buf_size = n_blocks * sizeof(block_q8_0); - block_q8_0 *quantized_data = (block_q8_0 *)malloc(n_blocks * sizeof(block_q8_0)); - if (!quantized_data) { + int64_t * buf_size = static_cast(buffer_size_msg.data()); + float * batch_embd = is_out_embd ? ubatch->out_embd : ubatch->backend_embd; + int64_t float_element_size = dims[0] * dims[1] * sizeof(float); + + quantized_array_t *quantized_array = (quantized_array_t*)malloc(*buf_size); + if (!quantized_array) { LLAMA_LOG_INFO("Failed to allocate space for recv data.\n"); return; } - std::memcpy(quantized_data, data_msg.data(), buf_size); - float * batch_embd = is_out_embd ? ubatch->out_embd : ubatch->backend_embd; - int64_t tensor_elements = dims[0] * dims[1]; - q8_0_dequantize(quantized_data, n_blocks, tensor_elements, batch_embd); - free(quantized_data); - - // size_t buf_size = dims[0] * dims[1] * sizeof(float); - // std::memcpy(batch_embd, data_msg.data(), buf_size); - - // TODO: Fix dump feature to support int8 - // if (dump_folder && strlen(dump_folder) > 0) { - // std::string dump_path = std::string(dump_folder) + "/recv_" + std::to_string(g_llama_recv_tensors_counts) + ".bin"; - // dump_tensors(dump_path, static_cast(TensorDataType::FLOAT32), - // static_cast(dims[0]), - // static_cast(dims[1]), - // buf_size, data_msg.data()); - // } + std::memcpy(quantized_array, data_msg.data(), *buf_size); + + dequantize(quantized_array, batch_embd); + free_quantized_array(quantized_array); + if (dump_folder && strlen(dump_folder) > 0) { + std::string dump_path = std::string(dump_folder) + "/recv_" + std::to_string(g_llama_recv_tensors_counts) + ".bin"; + dump_tensors(dump_path, static_cast(TensorDataType::FLOAT32), + static_cast(dims[0]), + static_cast(dims[1]), + float_element_size, batch_embd); + } } else if (key == "inp_pos") { int64_t * dims = static_cast(dims_msg.data()); size_t buf_size = dims[0] * sizeof(int32_t); diff --git a/src/quantization.cpp b/src/quantization.cpp new file mode 100644 index 00000000..de492a30 --- /dev/null +++ b/src/quantization.cpp @@ -0,0 +1,265 @@ +#include "quantization.h" + +static int64_t _get_q8_0_quantized_array_size(const quantized_array_t *quantized_array) { + if (!quantized_array) return 0; + return sizeof(uint8_t) /* quantized_type */ + + 3 * sizeof(uint64_t) /* num_elements, num_blocks, block_size */ + + quantized_array->num_blocks * sizeof(float) /* scales */ + + quantized_array->num_elements * sizeof(int8_t); /* data */ +} + +static int64_t _get_q4_0_quantized_array_size(const quantized_array_t *quantized_array) { + if (!quantized_array) return 0; + + const uint64_t num_elements_for_data = (quantized_array->num_elements + 1) / 2; + + return sizeof(quantized_array_t) /* quantized_type num_elements, num_blocks, block_size */ + + quantized_array->num_blocks * sizeof(float) /* scales */ + + num_elements_for_data * sizeof(int8_t); /* packed data */ +} + +int64_t get_quantized_array_size(const quantized_array_t *quantized_array) { + if (!quantized_array) return 0; + switch (quantized_array->quantized_type) { + case 0: + return _get_q8_0_quantized_array_size(quantized_array); + case 1: + return _get_q4_0_quantized_array_size(quantized_array); + default: + return 0; /* unknown type */ + } +} + +quantized_array_t *allocate_q8_0_array(uint64_t num_elements, + uint64_t block_size) +{ + if (!num_elements || !block_size) return NULL; + + uint64_t num_blocks = (num_elements + block_size - 1) / block_size; + + size_t total = sizeof(quantized_array_t) + + num_blocks * sizeof(float) + + num_elements * sizeof(int8_t); + + quantized_array_t *qa = (quantized_array_t*)calloc(1, total); + if (!qa) return NULL; + + /* initialise the header fields */ + qa->quantized_type = 0; /* q8_0 */ + qa->num_elements = num_elements; + qa->num_blocks = num_blocks; + qa->block_size = block_size; + + qa->scales = (float*)(qa + 1); /* just after the header */ + qa->data = (int8_t*)(qa->scales + num_blocks); /* after the scales */ + + return qa; +} + +quantized_array_t *allocate_q4_0_array(uint64_t num_elements, + uint64_t block_size) +{ + if (!num_elements || !block_size) return NULL; + + uint64_t num_blocks = (num_elements + block_size - 1) / block_size; + uint64_t num_elements_for_data = (num_elements + 1) / 2; + + size_t total = sizeof(quantized_array_t) + + num_blocks * sizeof(float) + + num_elements_for_data * sizeof(int8_t); + + quantized_array_t *qa = (quantized_array_t*)calloc(1, total); + if (!qa) return NULL; + + qa->quantized_type = 1; /* q4_0 */ + qa->num_elements = num_elements; + qa->num_blocks = num_blocks; + qa->block_size = block_size; + + qa->scales = (float*)(qa + 1); + qa->data = (int8_t*)(qa->scales + num_blocks); + + return qa; +} + +void free_quantized_array(quantized_array_t *quantized_array) { + if (!quantized_array) return; + free(quantized_array); +} + +static int _quantize_q8_0(const float *float_array, + quantized_array_t *quantized_array) { + if (!float_array || !quantized_array) return 1; + + const uint64_t block_size = quantized_array->block_size; + const uint64_t num_blocks = quantized_array->num_blocks; + const uint64_t num_elements = quantized_array->num_elements; + + for (uint64_t b = 0; b < num_blocks; ++b) { + const uint64_t start = b * block_size; + const uint64_t remain = (start + block_size <= num_elements) + ? block_size + : (num_elements - start); + + /* 1) find max‑abs in this block */ + float abs_max = 0.0f; + for (uint64_t i = 0; i < remain; ++i) { + float v = fabsf(float_array[start + i]); + if (v > abs_max) abs_max = v; + } + + /* 2) compute scale */ + float scale = (abs_max > 0.0f) ? (abs_max / 127.0f) : 0.0f; + float inv_scale = (scale > 0.0f) ? (1.0f / scale) : 0.0f; + quantized_array->scales[b] = scale; + + /* 3) quantise */ + for (uint64_t i = 0; i < remain; ++i) { + float val = float_array[start + i] * inv_scale; + long qi = lrintf(val); /* nearest int */ + if (qi < -127) qi = -127; + if (qi > 127) qi = 127; + quantized_array->data[start + i] = (int8_t)qi; + } + } + return 0; +} + +static int _quantize_q4_0(const float *float_array, + quantized_array_t *quantized_array) { + if (!float_array || !quantized_array) return 1; + + const uint64_t block_size = quantized_array->block_size; + const uint64_t num_blocks = quantized_array->num_blocks; + const uint64_t num_elements = quantized_array->num_elements; + uint8_t *data = (uint8_t *)quantized_array->data; + + for (uint64_t b = 0; b < num_blocks; ++b) { + const uint64_t start = b * block_size; + const uint64_t remain = (start + block_size <= num_elements) + ? block_size + : (num_elements - start); + + /* 1) find max‑abs in this block */ + float abs_max = 0.0f; + for (uint64_t i = 0; i < remain; ++i) { + float v = fabsf(float_array[start + i]); + if (v > abs_max) abs_max = v; + } + + /* 2) compute scale */ + float scale = (abs_max > 0.0f) ? (abs_max / 7.0f) : 0.0f; + float inv_scale = (scale > 0.0f) ? (1.0f / scale) : 0.0f; + quantized_array->scales[b] = scale; + + /* 3) quantise */ + for (uint64_t i = 0; i < remain; ++i) { + float val = float_array[start + i] * inv_scale; + long qi = lrintf(val); /* nearest int */ + if (qi < -7) qi = -7; + if (qi > 7) qi = 7; + + uint8_t four_bit_qi = ((uint8_t)qi) & 0x0F; + + int data_index = (start + i) / 2; + if (i % 2 == 0) { + data[data_index] = (uint8_t)(four_bit_qi << 4); + } + else { + data[data_index] = (uint8_t)(data[data_index] | four_bit_qi); + } + } + } + return 0; +} + +int quantize(const float *float_array, + uint64_t num_elements, + uint8_t quantized_type, + quantized_array_t **quantized_array) { + if (!float_array || num_elements == 0 || *quantized_array) return 1; + + switch (quantized_type) { + case 0: /* q8_0 */ + *quantized_array = allocate_q8_0_array(num_elements, + DEFAULT_Q8_0_BLOCK_SIZE); + if (!*quantized_array) return 1; + return _quantize_q8_0(float_array, *quantized_array); + + case 1: /* q4_0 */ + *quantized_array = allocate_q4_0_array(num_elements, + DEFAULT_Q4_0_BLOCK_SIZE); + if (!*quantized_array) return 1; + return _quantize_q4_0(float_array, *quantized_array); + default: + return 1; /* unknown type */ + } +} + +static int _dequantize_q8_0(const quantized_array_t *quantized_array, + float *float_array) { + const uint64_t block_size = quantized_array->block_size; + const uint64_t num_blocks = quantized_array->num_blocks; + const uint64_t num_elements = quantized_array->num_elements; + + for (uint64_t b = 0; b < num_blocks; ++b) { + const uint64_t start = b * block_size; + const uint64_t remain = (start + block_size <= num_elements) + ? block_size + : (num_elements - start); + const float scale = quantized_array->scales[b]; + + for (uint64_t i = 0; i < remain; ++i) { + float_array[start + i] = scale * (float)quantized_array->data[start + i]; + } + } + return 0; +} + +static int _dequantize_q4_0(const quantized_array_t *quantized_array, + float *float_array) { + const uint64_t block_size = quantized_array->block_size; + const uint64_t num_blocks = quantized_array->num_blocks; + const uint64_t num_elements = quantized_array->num_elements; + const uint8_t *src_data = (const uint8_t *)quantized_array->data; + + for (uint64_t b = 0; b < num_blocks; ++b) { + const uint64_t start = b * block_size; + const uint64_t remain = (start + block_size <= num_elements) + ? block_size + : (num_elements - start); + const float scale = quantized_array->scales[b]; + + for (uint64_t i = 0; i < remain; ++i) { + int data_index = (start + i) / 2; + uint8_t packed_qi = src_data[data_index]; + + if (i % 2 == 0) { + uint8_t qi = packed_qi >> 4; + int8_t signed_qi = (int8_t)(qi << 4) >> 4; + float_array[start + i] = scale * (float)(signed_qi); + } + else { + uint8_t qi = packed_qi & 0x0F; + int8_t signed_qi = (int8_t)(qi << 4) >> 4; + float_array[start + i] = scale * (float)(signed_qi); + } + } + } + return 0; +} + + +int dequantize(const quantized_array_t *quantized_array, float *float_array) { + if (!quantized_array || !float_array) return 1; + + switch (quantized_array->quantized_type) { + case 0: /* q8_0 */ + return _dequantize_q8_0(quantized_array, float_array); + + case 1: /* q4_0 */ + return _dequantize_q4_0(quantized_array, float_array); + default: + return 1; /* unknown type */ + } +} diff --git a/src/quantization.h b/src/quantization.h new file mode 100644 index 00000000..e832889d --- /dev/null +++ b/src/quantization.h @@ -0,0 +1,41 @@ +#ifndef QUANTIZATION_H +#define QUANTIZATION_H + +#include +#include +#include +#include + +/* The setting is refer to https://huggingface.co/docs/hub/en/gguf */ +#define DEFAULT_Q8_0_BLOCK_SIZE 32 +#define DEFAULT_Q4_0_BLOCK_SIZE 32 +#define DEFAULT_Q4_K_SUPER_BLOCK_SIZE 8 + +typedef struct { + uint8_t quantized_type; /* 0: q8_0, 1: q4_0, … */ + uint64_t num_elements; /* total elements in the original float array */ + uint64_t num_blocks; /* number of blocks (for block‑wised formats) */ + uint64_t block_size; /* elements per block */ + float *scales; /* length = num_blocks (or num_superblocks for kquant formats) */ + int8_t *data; /* for kquant, here need to contain quantized scale value + quantized value, otherwise it only need to store quantized value*/ +} quantized_array_t; + +quantized_array_t *allocate_q8_0_array(uint64_t num_elements, + uint64_t block_size); + +quantized_array_t *allocate_q4_0_array(uint64_t num_elements, + uint64_t block_size); + +void free_quantized_array(quantized_array_t *quantized_array); + +int64_t get_quantized_array_size(const quantized_array_t *quantized_array); + +int quantize(const float *float_array, + uint64_t num_elements, + uint8_t quantized_type, + quantized_array_t **quantized_array); + +int dequantize(const quantized_array_t *quantized_array, + float *float_array); + +#endif