From 73969b44b1148d04f6143dece47cdd85198f6525 Mon Sep 17 00:00:00 2001 From: DandinPower Date: Sat, 6 Dec 2025 21:08:10 +0800 Subject: [PATCH] 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) --- Makefile | 22 +- README.md | 12 + common/arg.cpp | 10 + common/common.cpp | 1 + common/common.h | 1 + docs/development/create_args_for_prima.md | 82 +++++++ include/bitsqueeze.h | 69 ++++++ include/llama.h | 1 + src/k_quantization.cpp | 218 ----------------- src/k_quantization.h | 47 ---- src/llama.cpp | 234 ++++++++---------- src/quantization.cpp | 281 ---------------------- src/quantization.h | 45 ---- src/sparsity.cpp | 124 ---------- src/sparsity.h | 37 --- 15 files changed, 282 insertions(+), 902 deletions(-) create mode 100644 docs/development/create_args_for_prima.md create mode 100644 include/bitsqueeze.h delete mode 100644 src/k_quantization.cpp delete mode 100644 src/k_quantization.h delete mode 100644 src/quantization.cpp delete mode 100644 src/quantization.h delete mode 100644 src/sparsity.cpp delete mode 100644 src/sparsity.h diff --git a/Makefile b/Makefile index 2350ee47..182cbd4f 100644 --- a/Makefile +++ b/Makefile @@ -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 \ diff --git a/README.md b/README.md index 2f891a8b..472dc5a5 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/common/arg.cpp b/common/arg.cpp index 1ca75faf..33c52a59 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -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()), diff --git a/common/common.cpp b/common/common.cpp index 9a06c271..ecba2019 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -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; diff --git a/common/common.h b/common/common.h index e5229ac4..8439afb7 100644 --- a/common/common.h +++ b/common/common.h @@ -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 diff --git a/docs/development/create_args_for_prima.md b/docs/development/create_args_for_prima.md new file mode 100644 index 00000000..d03ae154 --- /dev/null +++ b/docs/development/create_args_for_prima.md @@ -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.` 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) diff --git a/include/bitsqueeze.h b/include/bitsqueeze.h new file mode 100644 index 00000000..55e3a9e9 --- /dev/null +++ b/include/bitsqueeze.h @@ -0,0 +1,69 @@ +#ifndef BITSQUEEZE_H +#define BITSQUEEZE_H + +#include + +#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 diff --git a/include/llama.h b/include/llama.h index 478f3fa7..f52520fb 100644 --- a/include/llama.h +++ b/include/llama.h @@ -401,6 +401,7 @@ extern "C" { const char * comm_datatype; int comm_sparse_percentage; + int comm_compression_threshold; }; // model quantization parameters diff --git a/src/k_quantization.cpp b/src/k_quantization.cpp deleted file mode 100644 index 85e849c4..00000000 --- a/src/k_quantization.cpp +++ /dev/null @@ -1,218 +0,0 @@ -#include "k_quantization.h" -#include - -#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; -} \ No newline at end of file diff --git a/src/k_quantization.h b/src/k_quantization.h deleted file mode 100644 index b7043e76..00000000 --- a/src/k_quantization.h +++ /dev/null @@ -1,47 +0,0 @@ -#ifndef K_QUANTIZATION_H -#define K_QUANTIZATION_H - -#include -#include -#include -#include -#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 diff --git a/src/llama.cpp b/src/llama.cpp index 80947bbc..18dbf208 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -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 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(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; diff --git a/src/quantization.cpp b/src/quantization.cpp deleted file mode 100644 index 1217f2ab..00000000 --- a/src/quantization.cpp +++ /dev/null @@ -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 */ - } -} diff --git a/src/quantization.h b/src/quantization.h deleted file mode 100644 index 4efd09b1..00000000 --- a/src/quantization.h +++ /dev/null @@ -1,45 +0,0 @@ -#ifndef QUANTIZATION_H -#define QUANTIZATION_H - -/* To ensure can also compiled with c++ */ -#include -#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); - -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 diff --git a/src/sparsity.cpp b/src/sparsity.cpp deleted file mode 100644 index 18b1a515..00000000 --- a/src/sparsity.cpp +++ /dev/null @@ -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; -} diff --git a/src/sparsity.h b/src/sparsity.h deleted file mode 100644 index 3c3946f1..00000000 --- a/src/sparsity.h +++ /dev/null @@ -1,37 +0,0 @@ -#ifndef SPARSITY_H -#define SPARSITY_H - -#include -#include -#include -#include -#include - -/** - * @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