From b35e13d4f2e32f96b5d6d253b6f87cb6179da0d7 Mon Sep 17 00:00:00 2001 From: DandinPower Date: Wed, 24 Dec 2025 21:16:11 +0800 Subject: [PATCH] feat: add q2_k fast and modify q2_k logic, add preliminary support for rms_norm_aware_importance but currently assumption is wrong and will lead to corrupted result --- AGENTS.md | 5 ++ CHANGES.md | 2 +- common/arg.cpp | 2 +- include/bitsqueeze.h | 11 ++- src/llama.cpp | 160 ++++++++++++++++++++++++++++++++++++++++--- 5 files changed, 167 insertions(+), 13 deletions(-) create mode 100644 AGENTS.md diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 00000000..b86b547d --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,5 @@ +# Repository Guidelines + +## Fork Focus (prima.cpp) +- This fork prioritizes system optimization for distributed inference, especially networking paths (send/recv compression). +- Capture fork-specific behavior, new flags, and compatibility notes in `CHANGES.md`. \ No newline at end of file diff --git a/CHANGES.md b/CHANGES.md index e159d159..6f119b26 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -258,7 +258,7 @@ This feature allows each rank to specify the data type used when sending tensors A new CLI argument (`--comm_datatype TYPE`) sets the communication data type. Supported values: - f32 (default, no compression) -- integer based (q8_0, q4_0, q2_k, iq2_s, iq2_xs, iq2_xxs) +- integer based (q8_0, q4_0, q2_k, q2_k_fast, iq2_s, iq2_xs, iq2_xxs) - float based (f32, bf16, fp16, fp8, fp4, mxfp8, mxfp4, nvfp4, nf4, nf4_dq) - f32_sparsity (no quantization, but allows `--comm_sparse_percentage` to select the top-k features for each token based on the specified sparsity percentage) diff --git a/common/arg.cpp b/common/arg.cpp index eea9386a..4300164e 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2120,7 +2120,7 @@ gpt_params_context gpt_params_parser_init(gpt_params & params, llama_example ex, )); add_opt(llama_arg( {"--comm-datatype"}, "TYPE", - format("Datatype for communication, currently support f32, bf16, fp16, q8_0, q4_0, q2_k, iq2_s, iq2_xs, iq2_xxs, fp8, fp4, mxfp8, mxfp4, nvfp4, nf4, nf4_dq, or f32_sparsity (default: %s)", params.comm_datatype.c_str()), + format("Datatype for communication, currently support f32, bf16, fp16, q8_0, q4_0, q2_k, q2_k_fast, iq2_s, iq2_xs, iq2_xxs, fp8, fp4, mxfp8, mxfp4, nvfp4, nf4, nf4_dq, or f32_sparsity (default: %s)", params.comm_datatype.c_str()), [](gpt_params & params, const std::string & value) { params.comm_datatype = value; } diff --git a/include/bitsqueeze.h b/include/bitsqueeze.h index 55e3a9e9..ca9825fd 100644 --- a/include/bitsqueeze.h +++ b/include/bitsqueeze.h @@ -25,13 +25,15 @@ typedef enum { IQ2_XXS = 13, IQ2_XS = 14, IQ2_S = 15, + Q2_K_FAST = 16, + TOPK_IM = 17, } 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 */ + float sparse_ratio; /* only meaningful for TOPK, TOPK_IM */ } bsq_shape_t; typedef struct bitsqueeze_buffer { @@ -50,12 +52,17 @@ int bsq_compress_2d(const float *src, uint16_t num_features, float sparse_ratio, bsq_method_t method, - bitsqueeze_buffer_t **out); + bitsqueeze_buffer_t **out, + const float *im); int bsq_decompress(const bitsqueeze_buffer_t *buf, float *dst, uint64_t dst_num_elements); +int bsq_apply(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); diff --git a/src/llama.cpp b/src/llama.cpp index 9f3d3d59..896506c5 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -3001,6 +3001,9 @@ struct llama_model { llama_hparams hparams = {}; llama_vocab vocab; + // RMS Aware Add: Add next node first layer attn_norm ggml tensor + struct ggml_tensor * next_node_first_attn_norm = nullptr; + // TODO: should init all tensors to nullptr struct ggml_tensor * tok_embd = nullptr; struct ggml_tensor * type_embd = nullptr; @@ -3941,6 +3944,15 @@ static bool this_layer_is_mine( } } +// rms aware add: add a function to count next node first layer +static uint32_t get_next_node_first_layer_id(uint32_t n_world, uint32_t my_rank, const uint32_t * n_layer_window) { + uint32_t cumulative_layers = 0; + for (uint32_t start_rank=0; start_rank < my_rank; start_rank++){ + cumulative_layers += n_layer_window[start_rank]; + } + return cumulative_layers; +} + static int32_t map_layer_to_local_id( uint32_t layer_id, uint32_t n_world, @@ -7713,6 +7725,11 @@ static bool llm_load_tensors_impl( } } + // rms aware todo: remember to check flag + // rms aware add: calculate the + uint32_t next_node_first_layer_idx = 0; + next_node_first_layer_idx = get_next_node_first_layer_id(n_world, my_rank, n_layer_window); + // assign the input and output layers on CPU by default if (my_rank == 0) { model.buft_input = llama_default_buffer_type_cpu(model, true); @@ -7746,6 +7763,11 @@ static bool llm_load_tensors_impl( // for moe merged tensors ctx_size += ggml_tensor_overhead() * my_layers * 3; + // rms aware todo: need to check enable next_node_aware or not + // rms aware add: make sure it has cpu ctx + buft_layer_count[llama_default_buffer_type_cpu(model, true)]++; + ctx_size += ggml_tensor_overhead(); + std::map ctx_map; for (auto & it : buft_layer_count) { struct ggml_init_params params = { @@ -7800,6 +7822,22 @@ static bool llm_load_tensors_impl( model.layers.resize(my_layers); const auto tn = LLM_TN(model.arch); + // rms aware todo: need to check enable next_node_aware or not and ignore if the comm_datatype is f32 + // rms aware add: load tensor + if (model.arch != LLM_ARCH_QWEN2 and model.arch != LLM_ARCH_QWEN) { + throw std::runtime_error("Unsupported model arch for rms_aware_importance_matrix, currently only support Qwen Family"); + } + ggml_context * ctx_cpu = ctx_map.at(llama_default_buffer_type_cpu(model, true)); + model.next_node_first_attn_norm = ml.create_tensor(ctx_cpu, tn(LLM_TENSOR_ATTN_NORM, "weight", next_node_first_layer_idx), {n_embd}, 0, true); + if (model.next_node_first_attn_norm && + model.next_node_first_attn_norm->type != GGML_TYPE_F32) { + throw std::runtime_error(format( + "next_node_first_attn_norm has type %s, expected f32", + ggml_type_name(model.next_node_first_attn_norm->type))); + } + LLAMA_LOG_INFO("[rms_aware_importance_matrix] Successfully load next node first layer (layer_id=%d) attn normalization weight\n", next_node_first_layer_idx); + + switch (model.arch) { case LLM_ARCH_LLAMA: case LLM_ARCH_REFACT: @@ -18142,6 +18180,7 @@ static inline bsq_method_t convert_comm_datatype_string_to_enum(std::string comm 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 == "q2_k_fast") return Q2_K_FAST; else if (comm_datatype_string == "bf16") return BF16; else if (comm_datatype_string == "fp16") return FP16; else if (comm_datatype_string == "fp8") return FP8; @@ -18158,7 +18197,94 @@ static inline bsq_method_t convert_comm_datatype_string_to_enum(std::string comm 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) { +static void get_rms_norm_aware_importance_matrix( + const float *float_array, // a: [num_tokens * num_features] + const float *next_layer_attn_norm_weight, // b: [num_features] + uint16_t num_tokens, + uint16_t num_features, + float **importance_matrix) // out: [num_tokens * num_features] +{ + if (!importance_matrix) return; + *importance_matrix = NULL; + + if (!float_array || !next_layer_attn_norm_weight) return; + if (num_tokens == 0 || num_features == 0) return; + + size_t num_elements = (size_t)num_tokens * (size_t)num_features; + if (num_elements / (size_t)num_features != (size_t)num_tokens) return; // overflow guard + + float *temp_importance_matrix = (float *)calloc(num_elements, sizeof(float)); + if (!temp_importance_matrix) return; + + // Precompute b_k^2 once (b is shared across tokens). + float *b_sq = (float *)malloc((size_t)num_features * sizeof(float)); + if (!b_sq) { + free(temp_importance_matrix); + return; + } + for (uint16_t k = 0; k < num_features; ++k) { + float bk = next_layer_attn_norm_weight[k]; + b_sq[k] = bk * bk; + } + + // set this to match model's RMSNorm epsilon. + const float eps = 1e-6f; + + const float n = (float)num_features; + const float inv_n = 1.0f / n; + const float inv_n2 = inv_n * inv_n; + + for (uint16_t t = 0; t < num_tokens; ++t) { + uint32_t base = (uint32_t)t * (uint32_t)num_features; + + // Pass 1: compute r^2 and S_w for this token + float sum_a2 = 0.0f; + float S_w = 0.0f; // S_w = sum_i (b_i a_i)^2 = sum_i (b_i^2 a_i^2) + + for (uint16_t j = 0; j < num_features; ++j) { + float aj = float_array[base + j]; + float aj2 = aj * aj; + sum_a2 += aj2; + S_w += b_sq[j] * aj2; + } + + // r^2 = mean(a^2) (+ eps if matching real RMSNorm) + float r2 = sum_a2 * inv_n + eps; + if (r2 <= 0.0f) r2 = (eps > 0.0f) ? eps : 1e-12f; // safety fallback + + // We need 1/r^2, 1/r^4, 1/r^6 + float inv_r2 = 1.0f / r2; + float inv_r4 = inv_r2 * inv_r2; + float inv_r6 = inv_r4 * inv_r2; + + // Token-specific constants to reduce per-k work + float term2_factor = (-2.0f * inv_n) * inv_r4; // multiplies (b_k^2 * a_k^2) + float term3_factor = (inv_n2 * inv_r6) * S_w; // multiplies (a_k^2) + + // Pass 2: compute importance for each feature k + for (uint16_t k = 0; k < num_features; ++k) { + float ak = float_array[base + k]; + float ak2 = ak * ak; + float bk2 = b_sq[k]; + + // imp_k = b_k^2/r^2 - 2 b_k^2 a_k^2/(n r^4) + a_k^2 S_w/(n^2 r^6) + float imp = (bk2 * inv_r2) + + (bk2 * ak2 * term2_factor) + + (ak2 * term3_factor); + + // True value is >= 0, but float roundoff can produce tiny negatives. + if (imp < 0.0f) imp = 0.0f; + + // temp_importance_matrix[base + k] = imp; + temp_importance_matrix[base + k] = imp * ak2; + } + } + + free(b_sq); + *importance_matrix = temp_importance_matrix; +} + +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, const float * next_layer_attn_norm_weight = nullptr) { g_llama_send_tensors_counts++; try { std::vector send_msgs; @@ -18190,15 +18316,27 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba float sparse_ratio = (float) comm_sparse_percentage / 100; bitsqueeze_buffer_t *buf = NULL; + float *im_array = NULL; start_compute_time = get_iso8601_ms_timestamp(); - 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(); - - if (c_res || !buf) { - fprintf(stderr, "TOPK compress failed for array, ratio %.2f\n", sparse_ratio); - bsq_free(buf); - return ; + if (next_layer_attn_norm_weight != nullptr) { + // enable TOPK_IM + get_rms_norm_aware_importance_matrix(ubatch->backend_embd, next_layer_attn_norm_weight, tensors->sub_gf_out->ne[1], tensors->sub_gf_out->ne[0], &im_array); + int c_res = bsq_compress_2d(ubatch->backend_embd, tensors->sub_gf_out->ne[1], tensors->sub_gf_out->ne[0], sparse_ratio, TOPK_IM, &buf, im_array); + if (c_res || !buf) { + fprintf(stderr, "TOPK with importance matrix compress failed for array, ratio %.2f\n", sparse_ratio); + bsq_free(buf); + return ; + } } + else { + 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, NULL); + if (c_res || !buf) { + fprintf(stderr, "TOPK compress failed for array, ratio %.2f\n", sparse_ratio); + bsq_free(buf); + return ; + } + } + end_compute_time = get_iso8601_ms_timestamp(); buf_size = bsq_get_packed_size(buf); send_msgs.emplace_back("sub_gf_out", strlen("sub_gf_out")); @@ -18208,6 +18346,10 @@ static void llama_send_tensors(zmq::socket_t & socket, struct llama_ubatch * uba send_msgs.emplace_back(buf, buf_size); send_msgs.emplace_back(&buf_size, sizeof(buf_size)); + if (im_array) { + free(im_array); + } + 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()); @@ -18862,7 +19004,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, lctx.cparams.comm_compression_threshold); + 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, (float *) model.next_node_first_attn_norm->data); 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); }