feat: refactor and fix bug for q8_0 and can easily support q4_0 and fp32 in future

src/llama.cpp

@@ -18171,8 +18171,8 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
         free_quantized_array(quantized_array);
         
         if (enable_comm_compute_log) {
-            LLAMA_LOG_INFO("[%d][%s][compute][start][quantize send tensor]\n", my_rank, start_compute_time.c_str());
-            LLAMA_LOG_INFO("[%d][%s][compute][end][quantize send tensor]\n", my_rank, end_compute_time.c_str());
+            LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][quantize]\n", my_rank, start_compute_time.c_str());
+            LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][quantize]\n", my_rank, end_compute_time.c_str());
         }
 
         if (dump_folder && strlen(dump_folder) > 0) {
@@ -18207,24 +18207,24 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
             int64_t * dims       = static_cast<int64_t *>(dims_msg.data());
             int64_t * buf_size = static_cast<int64_t *>(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);
+            int64_t num_elements = dims[0] * dims[1];
+            int64_t float_element_size   = num_elements * sizeof(float);
 
-            quantized_array_t *quantized_array = (quantized_array_t*)malloc(*buf_size);
+            quantized_array_t *quantized_array = load_quantized_array_from_buffer(data_msg.data(), *buf_size);
             if (!quantized_array) {
-                LLAMA_LOG_INFO("Failed to allocate space for recv data.\n");   
+                LLAMA_LOG_INFO("Failed to load quantized array from buffer.\n");   
                 return;
             }
-            std::memcpy(quantized_array, data_msg.data(), *buf_size);
-            
+
             std::string start_compute_time = get_iso8601_ms_timestamp();
-            dequantize(quantized_array, batch_embd);
+            dequantize(quantized_array, batch_embd); 
             std::string end_compute_time = get_iso8601_ms_timestamp();
 
             free_quantized_array(quantized_array);
 
             if (enable_comm_compute_log) {
-                LLAMA_LOG_INFO("[%d][%s][compute][start][quantize send tensor]\n", my_rank, start_compute_time.c_str());
-                LLAMA_LOG_INFO("[%d][%s][compute][end][quantize send tensor]\n", my_rank, end_compute_time.c_str());
+                LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][dequantize]\n", my_rank, start_compute_time.c_str());
+                LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][dequantize]\n", my_rank, end_compute_time.c_str());
             }
             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";

src/quantization.cpp

@@ -2,8 +2,7 @@
 
 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 */
+    return sizeof(quantized_array_t)    /* quantized_type num_elements, num_blocks, block_size */
          + quantized_array->num_blocks * sizeof(float)   /* scales */
          + quantized_array->num_elements * sizeof(int8_t); /* data */
 }
@@ -31,8 +30,7 @@ int64_t get_quantized_array_size(const quantized_array_t *quantized_array) {
 }
 
 quantized_array_t *allocate_q8_0_array(uint64_t num_elements,
-                                       uint64_t block_size)
-{
+                                       uint64_t block_size) {
     if (!num_elements || !block_size) return NULL;
 
     uint64_t num_blocks = (num_elements + block_size - 1) / block_size;
@@ -57,8 +55,7 @@ quantized_array_t *allocate_q8_0_array(uint64_t num_elements,
 }
 
 quantized_array_t *allocate_q4_0_array(uint64_t num_elements,
-                                       uint64_t block_size)
-{
+                                       uint64_t block_size) {
     if (!num_elements || !block_size) return NULL;
 
     uint64_t num_blocks = (num_elements + block_size - 1) / block_size;
@@ -87,6 +84,25 @@ void free_quantized_array(quantized_array_t *quantized_array) {
     free(quantized_array);
 }
 
+quantized_array_t *load_quantized_array_from_buffer(const void *buffer, int64_t buffer_size) {
+    quantized_array_t *quantized_array = (quantized_array_t*)calloc(1, buffer_size);
+    if (!quantized_array) return NULL;
+    
+    std::memcpy(quantized_array, buffer, buffer_size);
+    switch (quantized_array->quantized_type) {
+        case 0: /* q8_0 */
+            quantized_array->scales = (float*)(quantized_array + 1);
+            quantized_array->data   = (int8_t*)(quantized_array->scales + quantized_array->num_blocks);
+            return quantized_array;
+        case 1: /* q4_0 */
+            quantized_array->scales = (float*)(quantized_array + 1);
+            quantized_array->data   = (int8_t*)(quantized_array->scales + quantized_array->num_blocks);
+            return quantized_array;
+        default: 
+            return NULL; /* unknown type */
+    }
+}
+
 static int _quantize_q8_0(const float *float_array,
                           quantized_array_t *quantized_array) {
     if (!float_array || !quantized_array) return 1;

src/quantization.h

@@ -1,6 +1,8 @@
 #ifndef QUANTIZATION_H
 #define QUANTIZATION_H
 
+/* To ensure can also compiled with c++ */
+#include <cstring> 
 #include <stdint.h>
 #include <stdlib.h>
 #include <string.h>
@@ -30,6 +32,8 @@ void free_quantized_array(quantized_array_t *quantized_array);
 
 int64_t get_quantized_array_size(const quantized_array_t *quantized_array);
 
+quantized_array_t *load_quantized_array_from_buffer(const void *buffer, int64_t buffer_size);
+
 int quantize(const float *float_array,
              uint64_t num_elements,
              uint8_t quantized_type,