electron-rare/prima-cpp
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: turn the communication datatype from float32 to int8

Browse Source
This commit is contained in:
DandinPower
2025-08-08 21:18:33 +08:00
parent 7480ffbd30
commit ff310e7882
1 changed files with 132 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama.cpp
+132 -24
View File
@@ -97,6 +97,84 @@
#include <string>
#include <iostream>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdio.h>
#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;
@@ -18131,29 +18209,42 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
g_llama_send_tensors_counts++;
try {
std::vector<zmq::message_t> send_msgs;
size_t buf_size = 0;
send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out"));
send_msgs.emplace_back("q8_0", strlen("q8_0"));
send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne));
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);
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<uint8_t>(TensorDataType::FLOAT32),
static_cast<uint64_t>(tensors->sub_gf_out->ne[0]),
static_cast<uint64_t>(tensors->sub_gf_out->ne[1]),
buf_size, ubatch->backend_embd);
}
// 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);
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);
send_msgs.emplace_back(0, sizeof(int64_t)); // extra frame that recv does not account for
send_msgs.emplace_back(ubatch->pos, buf_size);
}
zmq::send_multipart(socket, send_msgs);
free(quantized_data);
// 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<uint8_t>(TensorDataType::FLOAT32),
// static_cast<uint64_t>(tensors->sub_gf_out->ne[0]),
// static_cast<uint64_t>(tensors->sub_gf_out->ne[1]),
// buf_size, ubatch->backend_embd);
// }
} catch (const zmq::error_t& e) {
LLAMA_LOG_INFO("Failed to send tensor data: %s\n", e.what());
}
@@ -18167,24 +18258,41 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
return;
}
for (size_t i = 0; i < recv_msgs.size(); i += 3) {
for (size_t i = 0; i < recv_msgs.size(); i += 5) {
std::string key = recv_msgs[i].to_string();
zmq::message_t &dims_msg = recv_msgs[i + 1];
zmq::message_t &data_msg = recv_msgs[i + 2];
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];
if (key == "sub_gf_out") {
int64_t * dims = static_cast<int64_t *>(dims_msg.data());
size_t buf_size = dims[0] * dims[1] * sizeof(float);
float * batch_embd = is_out_embd ? ubatch->out_embd : ubatch->backend_embd;
std::memcpy(batch_embd, data_msg.data(), buf_size);
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<uint8_t>(TensorDataType::FLOAT32),
static_cast<uint64_t>(dims[0]),
static_cast<uint64_t>(dims[1]),
buf_size, data_msg.data());
int64_t * n_blocks_ptr = static_cast<int64_t*>(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) {
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<uint8_t>(TensorDataType::FLOAT32),
// static_cast<uint64_t>(dims[0]),
// static_cast<uint64_t>(dims[1]),
// buf_size, data_msg.data());
// }
} else if (key == "inp_pos") {
int64_t * dims = static_cast<int64_t *>(dims_msg.data());
size_t buf_size = dims[0] * sizeof(int32_t);