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

feat: use bitsqueeze as communication compression library to replace the previous quantization/k_quantization/sparsity. Also add a new argument call comm_compression_threshold to allow user can set a threshold for compression (good for only compress prefilling tokens)

Browse Source
This commit is contained in:
DandinPower
2025-12-06 21:08:10 +08:00
parent f90029e3c3
commit 73969b44b1
15 changed files with 282 additions and 902 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Makefile
+1 -21
View File
@@ -271,7 +271,7 @@ MK_CXXFLAGS = -std=c++11 -fPIC
MK_NVCCFLAGS = -std=c++11
MK_CPPFLAGS += -isystem /usr/local/include
MK_LDFLAGS += -L/usr/local/lib -lzmq
MK_LDFLAGS += -L/usr/local/lib -lzmq -lbitsqz
ifeq ($(UNAME_S),Darwin)
MK_CPPFLAGS += -isystem /opt/homebrew/include
@@ -955,9 +955,6 @@ OBJ_LLAMA = \
src/unicode.o \
src/unicode-data.o \
src/network-utils.o \
src/quantization.o \
src/sparsity.o \
src/k_quantization.o \
OBJ_COMMON = \
common/profiler.o \
@@ -1161,8 +1158,6 @@ src/llama.o: \
src/llama-sampling.h \
src/unicode.h \
src/network-utils.h \
src/quantization.h \
src/sparsity.h \
include/llama.h \
ggml/include/ggml-cuda.h \
ggml/include/ggml-metal.h \
@@ -1178,21 +1173,6 @@ src/llama-vocab.o: \
include/llama.h
$(CXX) $(CXXFLAGS) -c $< -o $@
src/quantization.o: \
src/quantization.cpp \
src/quantization.h
$(CXX) $(CXXFLAGS) -c $< -o $@
src/sparsity.o: \
src/sparsity.cpp \
src/sparsity.h
$(CXX) $(CXXFLAGS) -c $< -o $@ -fopenmp
src/k_quantization.o: \
src/k_quantization.cpp \
src/k_quantization.h
$(CXX) $(CXXFLAGS) -c $< -o $@
src/llama-grammar.o: \
src/llama-grammar.cpp \
src/llama-grammar.h \
README.md
+12
View File
@@ -118,6 +118,7 @@ Before using this project, ensure you have the following dependencies installed:
- zmq >= 4.3.2 (used for cross-device communication)
- HiGHS >= 1.9.0 (used for automatic workload distribution)
- CUDA (optional, if you have a GPU)
- BitSqueeze >= 0.1.1 (used for communication compression)
**Linux (e.g., Ubuntu):**
@@ -138,6 +139,17 @@ sudo make install
sudo ldconfig
```
For BitSqueeze, download and install from [source](https://github.com/DandinPower/BitSqueeze):
```bash
git clone https://github.com/DandinPower/BitSqueeze.git
cd BitSqueeze
cmake -B build_shared -DBUILD_SHARED_LIBS=ON -DCMAKE_BUILD_TYPE=Release
cmake --build build_shared --config Release
sudo cmake --install build_shared
sudo ldconfig
```
**macOS:**
```shell
common/arg.cpp
+10
View File
@@ -2132,6 +2132,16 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex,
params.comm_sparse_percentage = value;
}
));
add_opt(llama_arg(
{"--comm-compression-threshold"}, "N",
format("Minimum tensor token count required before applying communication compression; smaller tensors are sent as f32 regardless of comm_datatype (default: %d).", params.comm_compression_threshold),
[](gpt_params & params, int value) {
if (value < 0) {
throw std::invalid_argument("error: --comm-compression-threshold must be >= 0");
}
params.comm_compression_threshold = value;
}
));
add_opt(llama_arg(
{"--positive-file"}, "FNAME",
format("positive prompts file, one prompt per line (default: '%s')", params.cvector_positive_file.c_str()),
common/common.cpp
+1
View File
@@ -2116,6 +2116,7 @@ struct llama_context_params llama_context_params_from_gpt_params(const gpt_param
}
cparams.comm_sparse_percentage = params.comm_sparse_percentage;
cparams.comm_compression_threshold = params.comm_compression_threshold;
cparams.n_ctx = params.n_ctx;
cparams.n_predict = params.n_predict;
common/common.h
+1
View File
@@ -383,6 +383,7 @@ struct gpt_params {
std::string comm_datatype = "f32"; // data type for communication
int comm_sparse_percentage = 100;
int comm_compression_threshold = 0; // minimum element count before applying communication compression
};
// call once at the start of a program if it uses libcommon
docs/development/create_args_for_prima.md
+82
View File
@@ -0,0 +1,82 @@
# Adding New CLI Arguments for `llama_send_tensors` / `llama_recv_tensors`
This note walks through how to add new CLI arguments that flow into the distributed communication path in `src/llama.cpp`, including all touchpoints you must update.
## Key Data Flow
1. CLI parse → `gpt_params` (`common/arg.cpp`)
2. Copy into context params → `llama_context_params` (`common/common.cpp` via `llama_context_params_from_gpt_params`)
3. Propagate into runtime context → `lctx.cparams` (`src/llama.cpp` initialization)
4. Used by send/recv → `llama_send_tensors` / `llama_recv_tensors` (`src/llama.cpp`)
## Files to Update
- `common/arg.cpp`: register the CLI flag(s), set defaults, validate ranges.
- `common/common.h`: add fields to `gpt_params` (string/int/bool as needed).
- `common/common.cpp`:
- `llama_context_params_from_gpt_params`: copy new fields from `gpt_params` into `llama_context_params` (allocate/copy strings if needed).
- Ensure `llama_context_default_params` (in `src/llama.cpp`) sets sensible defaults.
- `include/llama.h` and `spm-headers/llama.h`: extend `struct llama_context_params` with the new field(s).
- `src/llama.cpp`:
- Thread new params through any places that consume them (e.g., `llama_send_tensors` / `llama_recv_tensors`).
- Validate inputs at use-site (range checks, supported values).
- If the param affects compression/sparsity/decompression, branch in `llama_send_tensors` and interpret tags in `llama_recv_tensors`.
## Step-by-Step Template
1. **Define parameter storage**
- Add to `struct gpt_params` in `common/common.h`.
- Set default values there.
2. **Expose via CLI**
- In `common/arg.cpp`, add a `llama_arg` entry:
- Flag names (short/long), help text, value hints.
- Handler writes into `gpt_params` (string/int/bool handler).
- Validate ranges here if you want early failure.
3. **Copy into runtime context**
- In `llama_context_params_from_gpt_params` (`common/common.cpp`):
- For strings, allocate and `strcpy` into `cparams`.
- For scalars, direct assignment.
4. **API surface**
- Add the field to `struct llama_context_params` in both headers:
- `include/llama.h`
- `spm-headers/llama.h`
- Update `llama_context_default_params` in `src/llama.cpp` with defaults (nullptr/0/false, or a literal default).
5. **Use in send/recv**
- `src/llama.cpp`:
- Accept the new argument in `llama_send_tensors` (and `llama_recv_tensors` if needed).
- Pass `lctx.cparams.<field>` when calling `llama_send_tensors` from the main decode loop.
- Implement behavior (e.g., choosing compression mode, sparse ratio, thresholds).
- Add validation guards near use; log or error out cleanly.
6. **(Optional) Docs/CHANGES**
- Document the new flag in `CHANGES.md` and any user-facing README if needed.
## Example: Existing `comm_datatype` / `comm_sparse_percentage`
- **CLI**: `--comm-datatype`, `--comm-sparse-percentage` (`common/arg.cpp`).
- **Params**: Stored in `gpt_params` (`common/common.h`) with defaults.
- **Context copy**: `llama_context_params_from_gpt_params` handles string allocation and scalar copy (`common/common.cpp`).
- **Headers**: `llama_context_params` exposes `const char * comm_datatype; int comm_sparse_percentage;` (`include/llama.h`, `spm-headers/llama.h`).
- **Defaults**: `llama_context_default_params` sets them to `nullptr` / 0 (`src/llama.cpp`).
- **Usage**: `llama_send_tensors` inspects `comm_datatype` and uses `comm_sparse_percentage` when `f32_sparsity` is selected; `llama_recv_tensors` reads the tag and auto-decompresses.
## Validation Tips
- For ranged ints, check on parse and on use; fail fast with a clear message.
- For enums/strings, normalize and validate against a small allowed list before hitting the hot path.
- Keep the wire format self-describing: include a datatype tag in the multipart messages so receivers can branch correctly.
## Quick checklist
- [ ] Field in `gpt_params` with default
- [ ] CLI flag in `common/arg.cpp`
- [ ] Copy into `llama_context_params_from_gpt_params`
- [ ] Field added to both `llama.h` headers
- [ ] Default set in `llama_context_default_params`
- [ ] Passed into `llama_send_tensors`/`llama_recv_tensors`
- [ ] Behavior implemented + input validation
- [ ] Docs updated (optional)
include/bitsqueeze.h
+69
View File
@@ -0,0 +1,69 @@
#ifndef BITSQUEEZE_H
#define BITSQUEEZE_H
#include <stdint.h>
#ifdef __cplusplus
extern "C" {
#endif
typedef enum {
BSQ_INVALID = -1,
Q8_0 = 0,
Q4_0 = 1,
Q2_K = 2,
TOPK = 3,
BF16 = 4,
FP16 = 5,
FP8 = 6,
FP4 = 7,
MXFP8 = 8,
MXFP4 = 9,
NVFP4 = 10,
NF4_DQ = 11,
NF4 = 12,
IQ2_XXS = 13,
IQ2_XS = 14,
IQ2_S = 15,
} bsq_method_t;
typedef struct {
uint64_t num_elements; /* for 1D formats */
uint16_t num_tokens; /* for 2D sparsity */
uint16_t num_features; /* for 2D sparsity */
float sparse_ratio; /* only meaningful for TOPK */
} bsq_shape_t;
typedef struct bitsqueeze_buffer {
bsq_method_t method;
bsq_shape_t shape;
void *payload;
} bitsqueeze_buffer_t;
int bsq_compress_1d(const float *src,
uint64_t num_elements,
bsq_method_t method,
bitsqueeze_buffer_t **out);
int bsq_compress_2d(const float *src,
uint16_t num_tokens,
uint16_t num_features,
float sparse_ratio,
bsq_method_t method,
bitsqueeze_buffer_t **out);
int bsq_decompress(const bitsqueeze_buffer_t *buf,
float *dst,
uint64_t dst_num_elements);
int64_t bsq_get_packed_size(const bitsqueeze_buffer_t *buf);
bitsqueeze_buffer_t *load_bsq_from_buffer(const void *buffer, int64_t buffer_size);
void bsq_free(bitsqueeze_buffer_t *buf);
#ifdef __cplusplus
}
#endif
#endif
include/llama.h
+1
View File
@@ -401,6 +401,7 @@ extern "C" {
const char * comm_datatype;
int comm_sparse_percentage;
int comm_compression_threshold;
};
// model quantization parameters
src/k_quantization.cpp
-218
View File
@@ -1,218 +0,0 @@
#include "k_quantization.h"
#include <stdio.h>
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
// The implementation is refer to https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-quants.c#L622
quantized_array_q2_k_t *allocate_q2_k_array(uint64_t num_elements) {
if (!num_elements) return NULL;
uint64_t num_elements_aligned = (num_elements % WEIGHT_PER_SUPER_BLOCK == 0) ? num_elements : num_elements + (WEIGHT_PER_SUPER_BLOCK - (num_elements % WEIGHT_PER_SUPER_BLOCK));
uint64_t num_super_blocks = num_elements_aligned / WEIGHT_PER_SUPER_BLOCK;
size_t total = sizeof(quantized_array_q2_k_t) + num_super_blocks * sizeof(super_block_q2_k);
quantized_array_q2_k_t *qa = (quantized_array_q2_k_t*)calloc(1, total);
if (!qa) return NULL;
qa->num_elements = num_elements;
qa->num_elements_aligned = num_elements_aligned;
qa->num_super_blocks = num_super_blocks;
qa->super_blocks = (super_block_q2_k*)(qa + 1);
return qa;
}
void free_quantized_q2_k_array(quantized_array_q2_k_t *quantized_array_q2_k) {
if (!quantized_array_q2_k) return;
free(quantized_array_q2_k);
}
int64_t get_quantized_q2_k_array_size(const quantized_array_q2_k_t *quantized_array_q2_k) {
if (!quantized_array_q2_k) return 0;
return sizeof(quantized_array_q2_k_t) + quantized_array_q2_k->num_super_blocks * sizeof(super_block_q2_k);
}
quantized_array_q2_k_t *load_quantized_q2_k_array_from_buffer(const void *buffer, int64_t buffer_size) {
quantized_array_q2_k_t *quantized_array = (quantized_array_q2_k_t*)calloc(1, buffer_size);
if (!quantized_array) return NULL;
memcpy(quantized_array, buffer, buffer_size);
quantized_array->super_blocks = (super_block_q2_k*)(quantized_array + 1);
return quantized_array;
}
static void find_optimal_scale_and_min(int curr_block_index, float *weights, float *scales, float*mins){
// naive approach
const float q2_scale = 3.f;
float min_val = INFINITY;
float max_val = -INFINITY;
for (int l = 0; l < Q2_K_BLOCK_SIZE; l++) {
if (weights[l] < min_val) min_val = weights[l];
}
for (int l = 0; l < Q2_K_BLOCK_SIZE; l++) {
weights[l] -= min_val;
}
for (int l = 0; l < Q2_K_BLOCK_SIZE; l++) {
if (weights[l] > max_val) max_val = weights[l];
}
scales[curr_block_index] = max_val / q2_scale;
mins[curr_block_index] = min_val;
}
int k_quantize(const float *float_array, uint64_t num_elements, quantized_array_q2_k_t **quantized_array_q2_k) {
const float q4_scale = 15.f;
uint8_t L[WEIGHT_PER_SUPER_BLOCK];
float weights[Q2_K_BLOCK_SIZE];
float mins[Q2_K_SUPER_BLOCK_SIZE];
float scales[Q2_K_SUPER_BLOCK_SIZE];
if (!float_array || num_elements == 0 || *quantized_array_q2_k) {
return 1;
}
*quantized_array_q2_k = allocate_q2_k_array(num_elements);
if (!*quantized_array_q2_k) {
return 1;
}
quantized_array_q2_k_t *qa = *quantized_array_q2_k;
float *float_array_aligned = (float*) calloc(1, sizeof(float) * qa->num_elements_aligned);
memcpy(float_array_aligned, float_array, (qa->num_elements) * sizeof(float));
for (uint32_t curr_super_block_index = 0; curr_super_block_index < qa->num_super_blocks; curr_super_block_index++) {
super_block_q2_k *curr_super_block = &qa->super_blocks[curr_super_block_index];
float max_scale = -INFINITY;
float max_abs_min = 0.f;
for (int j = 0; j < Q2_K_SUPER_BLOCK_SIZE; j++) {
for (int l = 0; l < Q2_K_BLOCK_SIZE; l++) {
weights[l] = float_array_aligned[j * Q2_K_BLOCK_SIZE + l];
}
find_optimal_scale_and_min(j, weights, scales, mins);
if (scales[j] > max_scale) {
max_scale = scales[j];
}
if (fabsf(mins[j]) > max_abs_min) {
max_abs_min = fabsf(mins[j]);
}
}
if (max_scale > 0) {
float iscale = q4_scale / max_scale;
for (int j = 0; j < Q2_K_SUPER_BLOCK_SIZE; j++) {
int l = (int)lrintf(iscale*scales[j]);
curr_super_block->scales[j] = l;
}
curr_super_block->super_scale = fp16_ieee_from_fp32_value(max_scale / q4_scale);
} else {
for (int j = 0; j < Q2_K_SUPER_BLOCK_SIZE; j++) curr_super_block->scales[j] = 0;
curr_super_block->super_scale = fp16_ieee_from_fp32_value(0.f);
}
if (max_abs_min > 0) {
const float iscale = 7.f / max_abs_min;
for (int j = 0; j < Q2_K_SUPER_BLOCK_SIZE; j++) {
int l = (int)lrintf(iscale * mins[j]);
l = MAX(-8, MIN(7, l));
curr_super_block->scales[j] |= ((l & 0xF) << 4);
}
curr_super_block->super_min = fp16_ieee_from_fp32_value(max_abs_min / 7.f);
} else {
curr_super_block->super_min = fp16_ieee_from_fp32_value(0.f);
}
for (int j = 0; j < Q2_K_SUPER_BLOCK_SIZE; j++) {
const float temp_scale = fp16_ieee_to_fp32_value(curr_super_block->super_scale) * (curr_super_block->scales[j] & 0xF);
const float m = fp16_ieee_to_fp32_value(curr_super_block->super_min);
const int8_t min_q = (curr_super_block->scales[j] >> 4);
const float temp_min = m * ((int8_t)(min_q << 4) >> 4);
for (int ii = 0; ii < Q2_K_BLOCK_SIZE; ii++) {
float val = (temp_scale > 0.f) ? (float_array_aligned[j * Q2_K_BLOCK_SIZE + ii] - temp_min) / temp_scale : 0.f;
int l = (int)lrintf(val);
l = MAX(0, MIN(3, l));
L[j * Q2_K_BLOCK_SIZE + ii] = l;
}
}
uint32_t packed_run = WEIGHT_PER_SUPER_BLOCK / 2; // 128
for (int j = 0; j < WEIGHT_PER_SUPER_BLOCK; j += packed_run) {
for (int l = 0; l < Q2_K_BLOCK_SIZE * 2; l++) { // l = 0..31
uint8_t b0 = L[j + l + 0];
uint8_t b1 = L[j + l + 32];
uint8_t b2 = L[j + l + 64];
uint8_t b3 = L[j + l + 96];
curr_super_block->data[j / 4 + l] = b0 | (b1 << 2) | (b2 << 4) | (b3 << 6);
}
}
float_array_aligned += WEIGHT_PER_SUPER_BLOCK;
}
return 0;
}
int k_dequantize(const quantized_array_q2_k_t *quantized_array_q2_k, float *float_array) {
if (!quantized_array_q2_k || !float_array || quantized_array_q2_k->num_super_blocks == 0) {
return 1;
}
for (uint32_t s = 0; s < quantized_array_q2_k->num_super_blocks; ++s) {
const super_block_q2_k *curr_super_block = &quantized_array_q2_k->super_blocks[s];
const float super_scale = fp16_ieee_to_fp32_value(curr_super_block->super_scale);
const float super_min = fp16_ieee_to_fp32_value(curr_super_block->super_min);
float scales[Q2_K_SUPER_BLOCK_SIZE];
float mins[Q2_K_SUPER_BLOCK_SIZE];
for(int i = 0; i < Q2_K_SUPER_BLOCK_SIZE; ++i) {
uint8_t packed_val = curr_super_block->scales[i];
scales[i] = super_scale * (packed_val & 0x0F);
int8_t min_q = (packed_val >> 4);
mins[i] = super_min * ((int8_t)(min_q << 4) >> 4);
}
const uint8_t *q = curr_super_block->data;
for (int l = 0; l < 32; ++l) {
uint8_t packed_byte = q[l];
int idx0 = l;
int idx1 = l + 32;
int idx2 = l + 64;
int idx3 = l + 96;
float_array[idx0] = mins[idx0/16] + scales[idx0/16] * ((packed_byte >> 0) & 3);
float_array[idx1] = mins[idx1/16] + scales[idx1/16] * ((packed_byte >> 2) & 3);
float_array[idx2] = mins[idx2/16] + scales[idx2/16] * ((packed_byte >> 4) & 3);
float_array[idx3] = mins[idx3/16] + scales[idx3/16] * ((packed_byte >> 6) & 3);
}
for (int l = 0; l < 32; ++l) {
uint8_t packed_byte = q[32 + l];
int idx0 = 128 + l;
int idx1 = 160 + l;
int idx2 = 192 + l;
int idx3 = 224 + l;
float_array[idx0] = mins[idx0/16] + scales[idx0/16] * ((packed_byte >> 0) & 3);
float_array[idx1] = mins[idx1/16] + scales[idx1/16] * ((packed_byte >> 2) & 3);
float_array[idx2] = mins[idx2/16] + scales[idx2/16] * ((packed_byte >> 4) & 3);
float_array[idx3] = mins[idx3/16] + scales[idx3/16] * ((packed_byte >> 6) & 3);
}
float_array += WEIGHT_PER_SUPER_BLOCK;
}
return 0;
}
src/k_quantization.h
-47
View File
@@ -1,47 +0,0 @@
#ifndef K_QUANTIZATION_H
#define K_QUANTIZATION_H
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include "fp16/fp16.h"
// The setting is refer to https://github.com/ggml-org/llama.cpp/blob/master/ggml/src/ggml-common.h
// fp16 implementation is refer to https://github.com/Maratyszcza/FP16/tree/master/include/fp16
#define Q2_K_BLOCK_SIZE 16
#define Q2_K_SUPER_BLOCK_SIZE 16
#define WEIGHT_PER_SUPER_BLOCK (Q2_K_BLOCK_SIZE*Q2_K_SUPER_BLOCK_SIZE)
// Q2_K 2-bit quantization
// weight is represented as x = a * q + b
// 16 blocks of 16 elements each
// 2.625 bits per weight ((16 * 4 * 2) + (256 * 2) + (16 * 2)) / 256 = 2.625
typedef struct {
uint16_t super_scale; // super-block scale for quantized scales (fp16)
uint16_t super_min; // super-block min for quantized scales (fp16)
uint8_t scales[Q2_K_SUPER_BLOCK_SIZE]; // scales and mins, quantized with 4 bits (length: Q2_K_SUPER_BLOCK_SIZE)
uint8_t data[WEIGHT_PER_SUPER_BLOCK / 4]; // quants with 2 bits (length: WEIGHT_PER_SUPER_BLOCK / 4)
} super_block_q2_k;
typedef struct {
uint64_t num_elements; /* total elements in the original float array */
uint64_t num_elements_aligned; /* aligned (padding) total elements for SUPER_BLOCK ELEMENTS */
uint32_t num_super_blocks;
super_block_q2_k *super_blocks;
} quantized_array_q2_k_t;
quantized_array_q2_k_t *allocate_q2_k_array(uint64_t num_elements);
void free_quantized_q2_k_array(quantized_array_q2_k_t *quantized_array_q2_k);
int64_t get_quantized_q2_k_array_size(const quantized_array_q2_k_t *quantized_array_q2_k);
quantized_array_q2_k_t *load_quantized_q2_k_array_from_buffer(const void *buffer, int64_t buffer_size);
int k_quantize(const float *float_array, uint64_t num_elements, quantized_array_q2_k_t **quantized_array_q2_k);
int k_dequantize(const quantized_array_q2_k_t *quantized_array_q2_k, float *float_array);
#endif
src/llama.cpp
+104 -128
View File
@@ -13,9 +13,7 @@
#include "profiler.h"
#include "network-utils.h"
#include "quantization.h"
#include "sparsity.h"
#include "k_quantization.h"
#include "bitsqueeze.h"
#ifdef GGML_USE_RPC
# include "ggml-rpc.h"
@@ -2721,6 +2719,7 @@ struct llama_cparams {
bool enable_comm_compute_log;
const char * comm_datatype;
float comm_sparse_percentage;
int comm_compression_threshold;
};
// TODO: separate into "llama_layer_enc" and "llama_layer_dec"
@@ -18139,14 +18138,38 @@ static int llama_recv_meta(zmq::socket_t & socket, struct sync_meta * meta) {
return 0;
}
static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, struct input_tensors * tensors, const char * dump_folder = nullptr, const bool enable_comm_compute_log = true, const int my_rank = 0, const char * comm_datatype = nullptr, int comm_sparse_percentage=100) {
static inline bsq_method_t convert_comm_datatype_string_to_enum(std::string comm_datatype_string) {
if (comm_datatype_string == "q8_0") return Q8_0;
else if (comm_datatype_string == "q4_0") return Q4_0;
else if (comm_datatype_string == "q2_k") return Q2_K;
else if (comm_datatype_string == "bf16") return BF16;
else if (comm_datatype_string == "fp16") return FP16;
else if (comm_datatype_string == "fp8") return FP8;
else if (comm_datatype_string == "fp4") return FP4;
else if (comm_datatype_string == "mxfp8") return MXFP8;
else if (comm_datatype_string == "mxfp4") return MXFP4;
else if (comm_datatype_string == "nvfp4") return NVFP4;
else if (comm_datatype_string == "nf4") return NF4;
else if (comm_datatype_string == "nf4_dq") return NF4_DQ;
else if (comm_datatype_string == "iq2_xxs") return IQ2_XXS;
else if (comm_datatype_string == "iq2_xs") return IQ2_XS;
else if (comm_datatype_string == "iq2_s") return IQ2_S;
else if (comm_datatype_string == "f32_sparsity") return TOPK;
else return BSQ_INVALID;
}
static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * ubatch, struct input_tensors * tensors, const char * dump_folder = nullptr, const bool enable_comm_compute_log = true, const int my_rank = 0, const char * comm_datatype = nullptr, int comm_sparse_percentage=100, int comm_compression_threshold=0) {
g_llama_send_tensors_counts++;
try {
std::vector<zmq::message_t> 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);
std::string comm_datatype_string = std::string(comm_datatype);
std::string comm_datatype_string = comm_datatype ? std::string(comm_datatype) : "f32";
int64_t compression_threshold = std::max<int64_t>(0, comm_compression_threshold);
if (tensors->sub_gf_out->ne[1] < compression_threshold) {
comm_datatype_string = "f32";
}
std::string start_compute_time = "";
std::string end_compute_time = "";
int64_t buf_size = 0;
@@ -18154,60 +18177,10 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
if (comm_datatype_string == "f32") {
buf_size = float_element_size;
send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out"));
send_msgs.emplace_back("normal", strlen("normal"));
send_msgs.emplace_back("f32", strlen("f32"));
send_msgs.emplace_back(tensors->sub_gf_out->ne, sizeof(tensors->sub_gf_out->ne));
send_msgs.emplace_back(ubatch->backend_embd, buf_size);
send_msgs.emplace_back(&buf_size, sizeof(int64_t));
} else if (comm_datatype_string == "q8_0" || comm_datatype_string == "q4_0") {
int qtype = (comm_datatype_string == "q8_0") ? 0 : 1;
start_compute_time = get_iso8601_ms_timestamp();
quantized_array_t *quantized_array = NULL;
if (quantize(ubatch->backend_embd, num_elements, qtype,
&quantized_array) || !quantized_array) {
LLAMA_LOG_INFO("Failed to allocate space or do quantization\n");
return;
}
end_compute_time = get_iso8601_ms_timestamp();
buf_size = get_quantized_array_size(quantized_array);
send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out"));
send_msgs.emplace_back("quantized", strlen("quantized"));
send_msgs.emplace_back(tensors->sub_gf_out->ne,
sizeof(tensors->sub_gf_out->ne));
send_msgs.emplace_back(quantized_array, buf_size);
send_msgs.emplace_back(&buf_size, sizeof(buf_size));
free_quantized_array(quantized_array);
if (enable_comm_compute_log) {
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());
}
} else if (comm_datatype_string == "q2_k") {
start_compute_time = get_iso8601_ms_timestamp();
quantized_array_q2_k_t *quantized_array = NULL;
if (k_quantize(ubatch->backend_embd, num_elements,
&quantized_array) || !quantized_array) {
LLAMA_LOG_INFO("Failed to allocate space or do quantization\n");
return;
}
end_compute_time = get_iso8601_ms_timestamp();
buf_size = get_quantized_q2_k_array_size(quantized_array);
send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out"));
send_msgs.emplace_back("k_quantized", strlen("k_quantized"));
send_msgs.emplace_back(tensors->sub_gf_out->ne,
sizeof(tensors->sub_gf_out->ne));
send_msgs.emplace_back(quantized_array, buf_size);
send_msgs.emplace_back(&buf_size, sizeof(buf_size));
free_quantized_q2_k_array(quantized_array);
if (enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][k_quantize]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][k_quantize]\n", my_rank, end_compute_time.c_str());
}
} else if (comm_datatype_string == "f32_sparsity") {
if (comm_sparse_percentage < 1 && comm_sparse_percentage > 100) {
fprintf(stderr, "Sparse percentage %d should between 1~100\n", comm_sparse_percentage);
@@ -18215,33 +18188,63 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
}
float sparse_ratio = (float) comm_sparse_percentage / 100;
sparse_array_t *sparse_array = NULL;
bitsqueeze_buffer_t *buf = NULL;
start_compute_time = get_iso8601_ms_timestamp();
if (compress(ubatch->backend_embd, tensors->sub_gf_out->ne[1], tensors->sub_gf_out->ne[0], sparse_ratio, &sparse_array)) {
fprintf(stderr, "compress failed for ratio %.2f\n", sparse_ratio);
free_sparse_array(sparse_array);
return;
}
int c_res = bsq_compress_2d(ubatch->backend_embd, tensors->sub_gf_out->ne[1], tensors->sub_gf_out->ne[0], sparse_ratio, TOPK, &buf);
end_compute_time = get_iso8601_ms_timestamp();
buf_size = get_sparse_array_size(sparse_array);
if (c_res || !buf) {
fprintf(stderr, "TOPK compress failed for array, ratio %.2f\n", sparse_ratio);
bsq_free(buf);
return ;
}
buf_size = bsq_get_packed_size(buf);
send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out"));
send_msgs.emplace_back("sparse", strlen("sparse"));
send_msgs.emplace_back("f32_sparsity", strlen("f32_sparsity"));
send_msgs.emplace_back(tensors->sub_gf_out->ne,
sizeof(tensors->sub_gf_out->ne));
send_msgs.emplace_back(sparse_array, buf_size);
send_msgs.emplace_back(buf, buf_size);
send_msgs.emplace_back(&buf_size, sizeof(buf_size));
free_sparse_array(sparse_array);
bsq_free(buf);
if (enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][sparse_compress]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][sparse_compress]\n", my_rank, end_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][compress]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][compress]\n", my_rank, end_compute_time.c_str());
}
} else {
LLAMA_LOG_INFO("Unsupported communication type = %s\n", comm_datatype_string);
return;
bsq_method_t bsq_method = convert_comm_datatype_string_to_enum(comm_datatype_string);
if (bsq_method != BSQ_INVALID) {
bitsqueeze_buffer_t *buf = NULL;
start_compute_time = get_iso8601_ms_timestamp();
int c_res = bsq_compress_1d(ubatch->backend_embd, num_elements, bsq_method, &buf);
end_compute_time = get_iso8601_ms_timestamp();
if (c_res || !buf) {
fprintf(stderr, "%s compress failed on array \n", comm_datatype_string.c_str());
bsq_free(buf);
return;
}
buf_size = bsq_get_packed_size(buf);
send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out"));
send_msgs.emplace_back(comm_datatype, strlen(comm_datatype));
send_msgs.emplace_back(tensors->sub_gf_out->ne,
sizeof(tensors->sub_gf_out->ne));
send_msgs.emplace_back(buf, buf_size);
send_msgs.emplace_back(&buf_size, sizeof(buf_size));
bsq_free(buf);
if (enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][compute][start][send_tensors][compress]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][send_tensors][compress]\n", my_rank, end_compute_time.c_str());
}
} else {
LLAMA_LOG_INFO("Unsupported communication type = %s\n", comm_datatype_string.c_str());
return;
}
}
if (tensors->inp_pos) {
@@ -18249,7 +18252,7 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
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("f32", strlen("f32"));
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));
@@ -18292,62 +18295,32 @@ static void llama_recv_tensors(zmq::socket_t & socket, struct llama_ubatch * uba
int64_t num_elements = dims[0] * dims[1];
int64_t float_element_size = num_elements * sizeof(float);
if (comm_type == "quantized") {
quantized_array_t *quantized_array = load_quantized_array_from_buffer(data_msg.data(), *buf_size);
if (!quantized_array) {
LLAMA_LOG_INFO("Failed to load quantized array from buffer.\n");
return;
}
std::string start_compute_time = get_iso8601_ms_timestamp();
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][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());
}
}
else if (comm_type == "k_quantized") {
quantized_array_q2_k_t *quantized_array = load_quantized_q2_k_array_from_buffer(data_msg.data(), *buf_size);
if (!quantized_array) {
LLAMA_LOG_INFO("Failed to load quantized array from buffer.\n");
return;
}
std::string start_compute_time = get_iso8601_ms_timestamp();
k_dequantize(quantized_array, batch_embd);
std::string end_compute_time = get_iso8601_ms_timestamp();
free_quantized_q2_k_array(quantized_array);
if (enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][k_dequantize]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][k_dequantize]\n", my_rank, end_compute_time.c_str());
}
}
else if (comm_type == "sparse") {
sparse_array_t *sparse_array = load_sparse_array_from_buffer(data_msg.data(), *buf_size);
if (!sparse_array) {
LLAMA_LOG_INFO("Failed to load sparse array from buffer.\n");
return;
}
std::string start_compute_time = get_iso8601_ms_timestamp();
decompress(sparse_array, batch_embd);
std::string end_compute_time = get_iso8601_ms_timestamp();
free_sparse_array(sparse_array);
if (enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][sparse_decompress]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][sparse_decompress]\n", my_rank, end_compute_time.c_str());
}
if (comm_type == "f32") {
std::memcpy(batch_embd, data_msg.data(), float_element_size);
}
else {
std::memcpy(batch_embd, data_msg.data(), float_element_size);
bsq_method_t bsq_method = convert_comm_datatype_string_to_enum(comm_type);
if (bsq_method != BSQ_INVALID) {
bitsqueeze_buffer_t *buf = load_bsq_from_buffer(data_msg.data(), *buf_size);
if (!buf) {
LLAMA_LOG_INFO("Failed to load bsq array from buffer.\n");
return;
}
std::string start_compute_time = get_iso8601_ms_timestamp();
bsq_decompress(buf, batch_embd, num_elements);
std::string end_compute_time = get_iso8601_ms_timestamp();
bsq_free(buf);
if (enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][compute][start][recv_tensors][decompress]\n", my_rank, start_compute_time.c_str());
LLAMA_LOG_INFO("[%d][%s][compute][end][recv_tensors][decompress]\n", my_rank, end_compute_time.c_str());
}
} else {
LLAMA_LOG_INFO("Unsupported communication type = %s\n", comm_type);
return;
}
}
if (dump_folder && strlen(dump_folder) > 0) {
@@ -18878,7 +18851,7 @@ static int llama_decode_internal(
struct input_tensors tensors = {sub_gf_out, lctx.inp_pos};
const bool is_to_master = my_rank != 0 && is_last_l;
zmq::socket_t * s = is_to_master ? lctx.master_socket : lctx.send_socket;
llama_send_tensors(*s, &ubatch, &tensors, lctx.cparams.dump_folder, lctx.cparams.enable_comm_compute_log, my_rank, lctx.cparams.comm_datatype, lctx.cparams.comm_sparse_percentage);
llama_send_tensors(*s, &ubatch, &tensors, lctx.cparams.dump_folder, lctx.cparams.enable_comm_compute_log, my_rank, lctx.cparams.comm_datatype, lctx.cparams.comm_sparse_percentage, lctx.cparams.comm_compression_threshold);
if (lctx.cparams.enable_comm_compute_log) {
LLAMA_LOG_INFO("[%d][%s][comm][end][send_tensors][sbatch_tokens: %lu, ubatch_tokens: %u, send the result to the next node or the master]\n", my_rank, get_iso8601_ms_timestamp().c_str(), lctx.sbatch.n_tokens, ubatch.n_tokens);
}
@@ -20638,7 +20611,9 @@ struct llama_context_params llama_context_default_params() {
/*.abort_callback_data =*/ nullptr,
/*.dump_folder =*/ nullptr,
/*.enable_comm_compute_log =*/ false,
/*.comm_datatype =*/ nullptr,
/*.comm_datatype =*/ nullptr,
/*.comm_sparse_percentage =*/ 100,
/*.comm_compression_threshold =*/ 0,
};
return result;
@@ -21285,6 +21260,7 @@ struct llama_context * llama_new_context_with_model(
ctx->cparams.enable_comm_compute_log = params.enable_comm_compute_log;
ctx->cparams.comm_datatype = params.comm_datatype;
ctx->cparams.comm_sparse_percentage = params.comm_sparse_percentage;
ctx->cparams.comm_compression_threshold = params.comm_compression_threshold;
ctx->cparams.original_next_rank = (params.rank + 1) % params.n_world;
auto &hparams = model->hparams;
src/quantization.cpp
-281
View File
@@ -1,281 +0,0 @@
#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(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 */
}
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);
}
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;
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 */
}
}
src/quantization.h
-45
View File
@@ -1,45 +0,0 @@
#ifndef QUANTIZATION_H
#define QUANTIZATION_H
/* To ensure can also compiled with c++ */
#include <cstring>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
/* 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);
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,
quantized_array_t **quantized_array);
int dequantize(const quantized_array_t *quantized_array,
float *float_array);
#endif
src/sparsity.cpp
-124
View File
@@ -1,124 +0,0 @@
#include "sparsity.h"
sparse_array_t *allocate_sparse_array(uint16_t num_tokens, uint16_t num_features, float sparse_ratio) {
if (!num_tokens || !num_features) return NULL;
if (sparse_ratio < 0.0f || sparse_ratio > 1.0f) return NULL;
float raw_sparse = (float)num_features * sparse_ratio;
uint16_t num_sparse_features = (uint16_t)roundf(raw_sparse);
// clamp to valid range
if (num_sparse_features > num_features) {
num_sparse_features = num_features;
} else if (num_sparse_features == 0 && sparse_ratio > 0.0f) {
num_sparse_features = 1; // Avoid total sparsity if ratio positive;
}
uint32_t sparse_elements = (uint32_t)num_tokens * num_sparse_features;
uint64_t total = sizeof(sparse_array_t) + sparse_elements * (sizeof(float) + sizeof(uint16_t));
sparse_array_t *sparse_array = (sparse_array_t*)calloc(1, total);
if (!sparse_array) return NULL;
/* initialise the header fields */
sparse_array->num_tokens = num_tokens;
sparse_array->num_features = num_features;
sparse_array->num_sparse_features = num_sparse_features;
sparse_array->sparse_indices = (uint16_t*)(sparse_array + 1); /* just after the header */
sparse_array->values = (float*)(sparse_array->sparse_indices + sparse_elements); /* after the sparse_indices */
return sparse_array;
}
void free_sparse_array(sparse_array_t *sparse_array) {
if (!sparse_array) return;
free(sparse_array);
}
uint64_t get_sparse_array_size(const sparse_array_t *sparse_array) {
if (!sparse_array) return 0;
uint32_t sparse_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_sparse_features;
return sizeof(sparse_array_t) + sparse_elements * (sizeof(float) + sizeof(uint16_t));
}
sparse_array_t *load_sparse_array_from_buffer(const void *buffer, uint64_t buffer_size) {
sparse_array_t *sparse_array = (sparse_array_t*)calloc(1, buffer_size);
if (!sparse_array) return NULL;
memcpy(sparse_array, buffer, buffer_size);
uint32_t sparse_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_sparse_features;
sparse_array->sparse_indices = (uint16_t*)(sparse_array + 1);
sparse_array->values = (float*)(sparse_array->sparse_indices + sparse_elements);
return sparse_array;
}
typedef struct {
uint16_t index;
float abs_val;
} sort_entry_t;
static int abs_sort_cmp(const void *a, const void *b) {
const float abs_a = ((const sort_entry_t *)a)->abs_val;
const float abs_b = ((const sort_entry_t *)b)->abs_val;
if (abs_a != abs_b) {
return (abs_a > abs_b) ? -1 : 1;
}
const uint16_t idx_a = ((const sort_entry_t *)a)->index;
const uint16_t idx_b = ((const sort_entry_t *)b)->index;
return (int)idx_a - (int)idx_b;
}
int compress(const float *float_array, uint16_t num_tokens, uint16_t num_features, float sparse_ratio, sparse_array_t **sparse_array) {
if (!float_array || num_tokens == 0 || num_features == 0 || *sparse_array) return 1;
/* ---- allocate sparse ------------------------------------------ */
*sparse_array = allocate_sparse_array(num_tokens, num_features, sparse_ratio);
if (!*sparse_array) return 1;
#pragma omp parallel for
for (uint16_t cur_token_index = 0; cur_token_index < num_tokens; cur_token_index++) {
sort_entry_t *entries = (sort_entry_t *)malloc(num_features * sizeof(sort_entry_t));
uint32_t dense_base = (uint32_t)cur_token_index * num_features;
uint32_t sparse_base = (uint32_t)cur_token_index * (*sparse_array)->num_sparse_features;
for (uint16_t i = 0; i < num_features; i++) {
entries[i].index = i;
entries[i].abs_val = fabsf(float_array[dense_base + i]);
}
qsort(entries, num_features, sizeof(sort_entry_t), abs_sort_cmp);
for (uint16_t keep_feature_index = 0; keep_feature_index < (*sparse_array)->num_sparse_features; keep_feature_index++) {
uint16_t orig_index = entries[keep_feature_index].index;
(*sparse_array)->sparse_indices[sparse_base + keep_feature_index] = orig_index;
(*sparse_array)->values[sparse_base + keep_feature_index] = float_array[dense_base + orig_index];
}
free(entries);
}
return 0;
}
int decompress(const sparse_array_t *sparse_array, float *float_array) {
if (!float_array || !sparse_array) return 1;
uint32_t num_elements = (uint32_t)sparse_array->num_tokens * sparse_array->num_features;
memset(float_array, 0, num_elements * sizeof(float));
for (uint16_t cur_token_index = 0; cur_token_index < sparse_array->num_tokens; cur_token_index++) {
uint32_t dense_base = (uint32_t)cur_token_index * sparse_array->num_features;
uint32_t sparse_base = (uint32_t)cur_token_index * sparse_array->num_sparse_features;
for (uint16_t keep_feature_index = 0; keep_feature_index < sparse_array->num_sparse_features; keep_feature_index++) {
uint16_t original_feature_index = sparse_array->sparse_indices[sparse_base + keep_feature_index];
float_array[dense_base + original_feature_index] = sparse_array->values[sparse_base + keep_feature_index];
}
}
return 0;
}
src/sparsity.h
-37
View File
@@ -1,37 +0,0 @@
#ifndef SPARSITY_H
#define SPARSITY_H
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <omp.h>
/**
* @brief Represents a sparse array in zero-based COO format for 2D data with shape [num_tokens, num_features].
*
* Sparsity is applied along the features dimension. Since each token retains the same number of sparse features,
* token indices are not stored explicitly. The structure holds the selected feature indices and corresponding values
* for all tokens in a flattened manner.
*/
typedef struct {
uint16_t num_tokens; /* Number of tokens (rows in the 2D shape). */
uint16_t num_features; /* Number of features per token (columns in the 2D shape). */
uint16_t num_sparse_features; /* Number of retained sparse features per token (must be <= num_features). */
uint16_t *sparse_indices; /* Flattened array of selected feature indices; length is (num_tokens * num_sparse_features). */
float *values; /* Flattened array of corresponding sparse values; length is (num_tokens * num_sparse_features). */
} sparse_array_t;
sparse_array_t *allocate_sparse_array(uint16_t num_tokens, uint16_t num_features, float sparse_ratio);
void free_sparse_array(sparse_array_t *sparse_array);
uint64_t get_sparse_array_size(const sparse_array_t *sparse_array);
sparse_array_t *load_sparse_array_from_buffer(const void *buffer, uint64_t buffer_size);
int compress(const float *float_array, uint16_t num_tokens, uint16_t num_features, float sparse_ratio, sparse_array_t **sparse_array);
int decompress(const sparse_array_t *sparse_array, float *float_array);
#endif