llama-cpp-bin 9509.0.0__tar.gz → 9518.0.0__tar.gz

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama-cpp-bin
-Version: 9509.0.0
+Version: 9518.0.0
 Summary: llama.cpp server binary built from source
 Author-email: Vladyslav Tkachenko <vladlearns@gmail.com>
 License: AGPL-3.0-or-later

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/llama.cpp/common/CMakeLists.txt

@@ -78,6 +78,8 @@ add_library(${TARGET}
     hf-cache.cpp
     hf-cache.h
     http.h
+    imatrix-loader.cpp
+    imatrix-loader.h
     json-partial.cpp
     json-partial.h
     json-schema-to-grammar.cpp

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/llama.cpp/common/arg.cpp

@@ -446,6 +446,12 @@ bool common_params_handle_models(common_params & params, llama_example curr_ex)
     opts.download_mtp    = spec_type_draft_mtp;
     opts.download_mmproj = !params.no_mmproj;
 
+    // sub-models (draft, mmproj, vocoder) are explicitly specified by the user,
+    // so we should not auto-discover mtp/mmproj siblings for them
+    common_download_opts sub_opts = opts;
+    sub_opts.download_mtp    = false;
+    sub_opts.download_mmproj = false;
+
     try {
         auto res = common_params_handle_model(params.model, opts);
         if (params.no_mmproj) {
@@ -457,7 +463,7 @@ bool common_params_handle_models(common_params & params, llama_example curr_ex)
         // only download mmproj if the current example is using it
         for (const auto & ex : mmproj_examples) {
             if (curr_ex == ex) {
-                common_params_handle_model(params.mmproj, opts);
+                common_params_handle_model(params.mmproj, sub_opts);
                 break;
             }
         }
@@ -470,8 +476,8 @@ bool common_params_handle_models(common_params & params, llama_example curr_ex)
             params.speculative.draft.mparams.url.empty()) {
             params.speculative.draft.mparams.path = res.mtp.path;
         }
-        common_params_handle_model(params.speculative.draft.mparams, opts);
-        common_params_handle_model(params.vocoder.model,             opts);
+        common_params_handle_model(params.speculative.draft.mparams, sub_opts);
+        common_params_handle_model(params.vocoder.model,             sub_opts);
         return true;
     } catch (const common_skip_download_exception &) {
         return false;

llama_cpp_bin-9518.0.0/llama.cpp/common/imatrix-loader.cpp

@@ -0,0 +1,165 @@
+#include "imatrix-loader.h"
+#include "common.h"
+#include "log.h"
+#include "gguf.h"
+
+#include <cmath>
+#include <cstring>
+#include <fstream>
+
+static bool common_imatrix_load_legacy(const std::string & fname, common_imatrix & imatrix) {
+    std::ifstream in(fname, std::ios::binary);
+    if (!in) {
+        LOG_ERR("%s: failed to open %s\n", __func__, fname.c_str());
+        return false;
+    }
+
+    int n_entries;
+    in.read((char *) &n_entries, sizeof(n_entries));
+    if (in.fail() || n_entries < 1) {
+        LOG_ERR("%s: no data in file %s\n", __func__, fname.c_str());
+        return false;
+    }
+
+    for (int i = 0; i < n_entries; ++i) {
+        int32_t len = 0;
+        in.read((char *) &len, sizeof(len));
+        std::vector<char> name_as_vec(len + 1);
+        in.read((char *) name_as_vec.data(), len);
+        if (in.fail()) {
+            LOG_ERR("%s: failed reading name for entry %d from %s\n", __func__, i + 1, fname.c_str());
+            return false;
+        }
+        name_as_vec[len] = 0;
+        std::string name{ name_as_vec.data() };
+
+        int32_t ncall = 0;
+        in.read((char *) &ncall, sizeof(ncall));
+        int32_t nval = 0;
+        in.read((char *) &nval, sizeof(nval));
+        if (in.fail() || nval < 1) {
+            LOG_ERR("%s: failed reading number of values for entry %d\n", __func__, i);
+            return false;
+        }
+
+        auto & e = imatrix.entries[std::move(name)];
+        e.sums.resize(nval);
+        in.read((char *) e.sums.data(), nval * sizeof(float));
+        if (in.fail()) {
+            LOG_ERR("%s: failed reading data for entry %d\n", __func__, i);
+            return false;
+        }
+
+        e.counts.resize(1);
+        e.counts[0] = ncall;
+    }
+
+    // the trailing data (chunk count + dataset name) is optional
+    if (in.peek() != EOF) {
+        int32_t n_calls = 0;
+        in.read((char *) &n_calls, sizeof(n_calls));
+        imatrix.chunk_count = n_calls;
+
+        if (!in.fail()) {
+            int32_t len = 0;
+            in.read((char *) &len, sizeof(len));
+            if (!in.fail() && len > 0) {
+                std::vector<char> dataset(len + 1, 0);
+                in.read(dataset.data(), len);
+                if (!in.fail()) {
+                    imatrix.datasets.push_back(dataset.data());
+                }
+            }
+        }
+    }
+
+    imatrix.chunk_size = 0;
+    imatrix.is_legacy  = true;
+
+    return true;
+}
+
+bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix) {
+    struct ggml_context * ctx = nullptr;
+    struct gguf_init_params meta_gguf_params = {
+        /* .no_alloc = */ false,
+        /* .ctx      = */ &ctx,
+    };
+    struct gguf_context * ctx_gguf = gguf_init_from_file(fname.c_str(), meta_gguf_params);
+    if (!ctx_gguf) {
+        return common_imatrix_load_legacy(fname, imatrix);
+    }
+
+    const int32_t n_entries = gguf_get_n_tensors(ctx_gguf);
+    if (n_entries < 1) {
+        LOG_ERR("%s: no data in file %s\n", __func__, fname.c_str());
+        gguf_free(ctx_gguf);
+        ggml_free(ctx);
+        return false;
+    }
+
+    const int64_t datasets_key   = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_DATASETS);
+    const int64_t chunk_count_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_COUNT);
+    const int64_t chunk_size_key  = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_CHUNK_SIZE);
+
+    if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
+        const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key);
+        imatrix.datasets.reserve(imatrix.datasets.size() + n);
+        for (int64_t i = 0; i < n; ++i) {
+            imatrix.datasets.push_back(gguf_get_arr_str(ctx_gguf, datasets_key, i));
+        }
+    }
+
+    imatrix.has_metadata = (datasets_key != -1 && chunk_count_key != -1 && chunk_size_key != -1);
+    imatrix.chunk_count  = (chunk_count_key != -1) ? gguf_get_val_u32(ctx_gguf, chunk_count_key) : 0;
+    imatrix.chunk_size   = (chunk_size_key  != -1) ? gguf_get_val_u32(ctx_gguf, chunk_size_key)  : 0;
+
+    const std::string in_sum2_suffix{ ".in_sum2" };
+    const std::string counts_suffix{ ".counts" };
+
+    std::map<std::string, std::pair<struct ggml_tensor *, struct ggml_tensor *>> sums_counts_for;
+
+    for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
+        std::string name = cur->name;
+
+        if (name.empty()) { continue; }
+
+        if (string_remove_suffix(name, in_sum2_suffix)) {
+            sums_counts_for[std::move(name)].first = cur;
+        } else if (string_remove_suffix(name, counts_suffix)) {
+            sums_counts_for[std::move(name)].second = cur;
+        }
+    }
+
+    for (const auto & sc : sums_counts_for) {
+        const std::string &        name    = sc.first;
+        const struct ggml_tensor * in_sum2 = sc.second.first;
+        const struct ggml_tensor * counts  = sc.second.second;
+
+        if (!in_sum2 || !counts) {
+            LOG_ERR("%s: mismatched sums and counts for %s\n", __func__, name.c_str());
+            gguf_free(ctx_gguf);
+            ggml_free(ctx);
+            return false;
+        }
+
+        auto & e = imatrix.entries[name];
+
+        const int64_t nval    = ggml_nelements(in_sum2);
+        const int64_t ncounts = ggml_nelements(counts);
+
+        e.sums.resize(nval);
+        for (int64_t j = 0; j < nval; ++j) {
+            e.sums[j] = ((const float *) in_sum2->data)[j];
+        }
+
+        e.counts.resize(ncounts);
+        for (int64_t j = 0; j < ncounts; ++j) {
+            e.counts[j] = std::lround(((const float *) counts->data)[j]);
+        }
+    }
+
+    gguf_free(ctx_gguf);
+    ggml_free(ctx);
+    return true;
+}

llama_cpp_bin-9518.0.0/llama.cpp/common/imatrix-loader.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <cstdint>
+#include <map>
+#include <string>
+#include <vector>
+
+inline constexpr const char * LLM_KV_IMATRIX_DATASETS    = "imatrix.datasets";
+inline constexpr const char * LLM_KV_IMATRIX_CHUNK_COUNT = "imatrix.chunk_count";
+inline constexpr const char * LLM_KV_IMATRIX_CHUNK_SIZE  = "imatrix.chunk_size";
+
+struct common_imatrix_entry {
+    std::vector<float>   sums;
+    std::vector<int64_t> counts;
+};
+
+struct common_imatrix {
+    std::map<std::string, common_imatrix_entry> entries;
+    std::vector<std::string> datasets;
+    int32_t chunk_count    = 0;
+    int32_t chunk_size     = 0;
+    bool    is_legacy      = false;
+    bool    has_metadata   = false;
+};
+
+bool common_imatrix_load(const std::string & fname, common_imatrix & imatrix);

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/llama.cpp/examples/speculative-simple/speculative-simple.cpp

@@ -175,7 +175,7 @@ int main(int argc, char ** argv) {
                     llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), seq_id));
 
             if (use_ckpt_dft) {
-                ckpt.update_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+                ckpt.update_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
             }
 
             // generate a new draft
@@ -196,12 +196,12 @@ int main(int argc, char ** argv) {
             // this allows us to restore the state if partial draft acceptance occurs
             if (!draft.empty()) {
                 if (use_ckpt_tgt) {
-                    ckpt.update_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+                    ckpt.update_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
                 }
             }
 
             {
-                ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+                ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
 
                 llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1);
             }
@@ -261,13 +261,13 @@ int main(int argc, char ** argv) {
             draft = std::move(ids);
 
             {
-                ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+                ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
 
                 llama_memory_seq_rm(llama_get_memory(ctx_tgt), seq_id, ckpt.pos_max + 1, -1);
             }
 
             {
-                ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+                ckpt.load_dft(ctx_dft.get(), seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
 
                 llama_memory_seq_rm(llama_get_memory(ctx_dft.get()), seq_id, ckpt.pos_max + 1, -1);
             }

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/llama.cpp/ggml/src/ggml-cpu/arch/wasm/quants.c

@@ -355,6 +355,78 @@ void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
     *s = sumf;
 }
 
+void ggml_vec_dot_q4_1_q8_1(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK8_1;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q4_1 * GGML_RESTRICT x = vx;
+    const block_q8_1 * GGML_RESTRICT y = vy;
+
+    float sumf = 0;
+
+#if defined __wasm_simd128__
+    v128_t sumv = wasm_f32x4_splat(0.0f);
+    float summs = 0.0f;
+
+    for (int ib = 0; ib < nb; ++ib) {
+        const block_q4_1 * GGML_RESTRICT x0 = &x[ib];
+        const block_q8_1 * GGML_RESTRICT y0 = &y[ib];
+
+        summs += GGML_CPU_FP16_TO_FP32(x0->m) * GGML_CPU_FP16_TO_FP32(y0->s);
+
+        const v128_t raw  = wasm_v128_load(x0->qs);
+        const v128_t v0s  = wasm_v128_and(raw, wasm_i8x16_splat(0x0F));
+        const v128_t v1s  = wasm_u8x16_shr(raw, 4);
+
+        const v128_t ys_lo = wasm_v128_load(y0->qs);
+        const v128_t ys_hi = wasm_v128_load(y0->qs + 16);
+
+        const v128_t v0s_l = wasm_u16x8_extend_low_u8x16(v0s);
+        const v128_t v0s_h = wasm_u16x8_extend_high_u8x16(v0s);
+        const v128_t ylo_l = wasm_i16x8_extend_low_i8x16(ys_lo);
+        const v128_t ylo_h = wasm_i16x8_extend_high_i8x16(ys_lo);
+        const v128_t v1s_l = wasm_u16x8_extend_low_u8x16(v1s);
+        const v128_t v1s_h = wasm_u16x8_extend_high_u8x16(v1s);
+        const v128_t yhi_l = wasm_i16x8_extend_low_i8x16(ys_hi);
+        const v128_t yhi_h = wasm_i16x8_extend_high_i8x16(ys_hi);
+
+        const v128_t acc = wasm_i32x4_add(
+            wasm_i32x4_add(
+                wasm_i32x4_dot_i16x8(v0s_l, ylo_l),
+                wasm_i32x4_dot_i16x8(v0s_h, ylo_h)),
+            wasm_i32x4_add(
+                wasm_i32x4_dot_i16x8(v1s_l, yhi_l),
+                wasm_i32x4_dot_i16x8(v1s_h, yhi_h)));
+
+        sumv = wasm_f32x4_add(sumv,
+            wasm_f32x4_mul(
+                wasm_f32x4_convert_i32x4(acc),
+                wasm_f32x4_splat(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d))));
+    }
+
+    sumf = wasm_f32x4_extract_lane(sumv, 0) + wasm_f32x4_extract_lane(sumv, 1) +
+           wasm_f32x4_extract_lane(sumv, 2) + wasm_f32x4_extract_lane(sumv, 3) + summs;
+
+    *s = sumf;
+
+#else
+    UNUSED(nb);
+    UNUSED(x);
+    UNUSED(y);
+    UNUSED(sumf);
+
+    ggml_vec_dot_q4_1_q8_1_generic(
+        n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
 void ggml_vec_dot_q5_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;
     const int nb = n / qk;

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/llama.cpp/src/llama-model.cpp

@@ -2112,6 +2112,15 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
                         filter = [n_main](int32_t il) { return (uint32_t)il >= n_main; };
                     }
 
+                    if (arch == LLM_ARCH_STEP35 && hparams.nextn_predict_layers > 0) {
+                        const uint32_t n_main = hparams.n_layer - hparams.nextn_predict_layers;
+                        if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
+                            filter = [n_main](int32_t il) { return (uint32_t)il >= n_main; };
+                        } else {
+                            filter = [n_main](int32_t il) { return (uint32_t)il <  n_main; };
+                        }
+                    }
+
                     if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) {
                         GGML_ASSERT(hparams.is_swa_any());
 

{llama_cpp_bin-9509.0.0 → llama_cpp_bin-9518.0.0}/llama.cpp/tools/imatrix/imatrix.cpp

@@ -1,5 +1,6 @@
 #include "arg.h"
 #include "common.h"
+#include "imatrix-loader.h"
 #include "log.h"
 #include "llama.h"
 #include "gguf.h"
@@ -34,10 +35,6 @@ static void print_usage(int, char ** argv) {
     LOG("\n");
 }
 
-static const char * const LLM_KV_IMATRIX_DATASETS    = "imatrix.datasets";
-static const char * const LLM_KV_IMATRIX_CHUNK_COUNT = "imatrix.chunk_count";
-static const char * const LLM_KV_IMATRIX_CHUNK_SIZE  = "imatrix.chunk_size";
-
 struct Stats {
     std::vector<float>   values;
     std::vector<int64_t> counts;
@@ -65,7 +62,6 @@ public:
     bool collect_imatrix(struct ggml_tensor * t, bool ask, void * user_data);
     void save_imatrix_legacy(int32_t ncall = -1) const;
     void save_imatrix(int32_t n_chunk = -1) const;
-    bool load_imatrix_legacy(const char * fname);
     bool load_imatrix(const char * file_name);
     const std::unordered_map<std::string, Stats> & get_mstats() const { return m_stats; }
 private:
@@ -624,204 +620,63 @@ void IMatrixCollector::save_imatrix(int32_t n_chunk) const {
     ggml_free(ctx);
 }
 
-bool IMatrixCollector::load_imatrix_legacy(const char * fname) {
-    std::ifstream in(fname, std::ios::binary);
-    if (!in) {
-        LOG_ERR("%s: failed to open %s\n", __func__, fname);
-        return false;
-    }
-    int n_entries;
-    in.read((char *) &n_entries, sizeof(n_entries));
-    if (in.fail() || n_entries < 1) {
-        LOG_ERR("%s: no data in file %s\n", __func__, fname);
+bool IMatrixCollector::load_imatrix(const char * file_name) {
+    common_imatrix loaded;
+    if (!common_imatrix_load(file_name, loaded)) {
         return false;
     }
-    // Guess the chunk size because it's not stored in the file
-    const int32_t chunk_size = m_params.n_ctx / m_params.n_parallel;
-
-    for (int i = 0; i < n_entries; ++i) {
-        int32_t len = 0;
-        in.read((char *) &len, sizeof(len));
-        std::vector<char> name_as_vec(len + 1);
-        in.read((char *) name_as_vec.data(), len);
-        if (in.fail()) {
-            LOG_ERR("%s: failed reading name for entry %d from %s\n", __func__, i + 1, fname);
-            return false;
-        }
-        name_as_vec[len] = 0;
-        std::string name{ name_as_vec.data() };
-        auto & e = m_stats[std::move(name)];
-        int32_t ncall = 0;
-        in.read((char *) &ncall, sizeof(ncall));
-        int32_t nval = 0;
-        in.read((char *) &nval, sizeof(nval));
-        if (in.fail() || nval < 1) {
-            LOG_ERR("%s: failed reading number of values for entry %d\n", __func__, i);
-            m_stats = {};
-            return false;
-        }
 
-        if (e.values.empty()) {
-            e.values.resize(nval, 0.0f);
-            e.counts.resize(1, 0);
-        }
-
-        std::vector<float> tmp(nval);
-        in.read((char *) tmp.data(), nval * sizeof(float));
-        if (in.fail()) {
-            LOG_ERR("%s: failed reading data for entry %d\n", __func__, i);
-            m_stats = {};
-            return false;
-        }
+    const int32_t chunk_size = m_params.n_ctx / m_params.n_parallel;
+    const bool is_legacy = loaded.is_legacy;
 
-        // Recreate the state as expected by save_imatrix(), and correct for weighted sum.
-        for (int i = 0; i < nval; i++) {
-            e.values[i] += tmp[i] * chunk_size;
-        }
-        // The legacy format doesn't distinguish the counts for different experts
-        for (size_t j = 0; j < e.counts.size(); ++j) {
-            e.counts[j] += ncall * chunk_size;
-        }
-    }
+    for (auto & [name, entry] : loaded.entries) {
+        auto & e = m_stats[name];
 
-    {
-        // TODO: extract into its own method; this is also used by the GGUF-based format
-        // Calculate the last chunk count
-        int64_t max_count = 0;
-        for (const auto & stats : m_stats) {
-            for (int64_t count : stats.second.counts) {
-                if (count > max_count) {
-                    max_count = count;
-                }
+        if (is_legacy) {
+            // Legacy format: sums contain (raw_sum/raw_count)*ncall, counts contain {ncall}
+            // Reconstruct raw form by multiplying by chunk_size
+            if (e.values.empty()) {
+                e.values.resize(entry.sums.size(), 0.0f);
+                e.counts.resize(1, 0);
             }
-        }
-        m_last_chunk = max_count / (chunk_size);
-    }
-
-    {
-        // Read the number of calls the matrix was computed with
-        int32_t n_calls;
-        in.read((char *) &n_calls, sizeof(n_calls));
-        // ignore it because it's not important
-    }
-
-    // Read the dataset path to include it when writing to GGUF
-    if (!in.fail()){
-        int32_t len = 0;
-        in.read((char *) &len, sizeof(len));
-        if (!in.fail()) {
-            std::vector<char> dataset;
-            dataset.resize(len + 1, 0);
-            in.read(dataset.data(), len);
-            if (!in.fail()) {
-                m_datasets.push_back(dataset.data());
+            for (size_t j = 0; j < entry.sums.size(); ++j) {
+                e.values[j] += entry.sums[j] * chunk_size;
+            }
+            for (size_t j = 0; j < e.counts.size(); ++j) {
+                e.counts[j] += entry.counts[0] * chunk_size;
             }
-        }
-    }
-
-    return true;
-}
-
-// Using GGUF as the file format, for greater extensibility
-bool IMatrixCollector::load_imatrix(const char * file_name) {
-    struct ggml_context * ctx = nullptr;
-    struct gguf_init_params meta_gguf_params = {
-        /* .no_alloc = */ false, // the data is needed
-        /* .ctx      = */ &ctx,
-    };
-    struct gguf_context * ctx_gguf = gguf_init_from_file(file_name, meta_gguf_params);
-    if (!ctx_gguf) {
-        return this->load_imatrix_legacy(file_name);
-    }
-    const int32_t n_entries = gguf_get_n_tensors(ctx_gguf);
-    if (n_entries < 1) {
-        LOG_ERR("%s: no data in file %s\n", __func__, file_name);
-        gguf_free(ctx_gguf);
-        ggml_free(ctx);
-        return false;
-    }
-
-    const int64_t datasets_key = gguf_find_key(ctx_gguf, LLM_KV_IMATRIX_DATASETS);
-    if (datasets_key != -1 && gguf_get_arr_type(ctx_gguf, datasets_key) == GGUF_TYPE_STRING) {
-        const int64_t n = gguf_get_arr_n(ctx_gguf, datasets_key);
-        m_datasets.reserve(m_datasets.size() + n);
-        for (int64_t i = 0; i < n; ++i) {
-            m_datasets.push_back(gguf_get_arr_str(ctx_gguf, datasets_key, i));
-        }
-    }
-
-    const std::string in_sum2_suffix{ ".in_sum2" };
-    const std::string counts_suffix{ ".counts" };
-
-    // Could re-use m_stats instead, but this allows
-    // checking for completeness of *each* loaded imatrix file
-    // and also makes it easier to re-use a similar implementation in quantize.cpp
-    // Using an ordered map to get a deterministic iteration order.
-    std::map<std::string, std::pair<struct ggml_tensor *, struct ggml_tensor *>> sums_counts_for;
-
-    for (struct ggml_tensor * cur = ggml_get_first_tensor(ctx); cur; cur = ggml_get_next_tensor(ctx, cur)) {
-        std::string name = cur->name;
-
-        if (name.empty()) { continue; }
-
-        if (string_remove_suffix(name, in_sum2_suffix)) {
-            // in_sum2
-            sums_counts_for[std::move(name)].first = cur;
-        } else if (string_remove_suffix(name, counts_suffix)) {
-            // counts
-            sums_counts_for[std::move(name)].second = cur;
         } else {
-            // ignore other tensors
-        }
-    }
-
-    for (const auto & sc : sums_counts_for) {
-        const std::string &        name    = sc.first;
-        const struct ggml_tensor * in_sum2 = sc.second.first;
-        const struct ggml_tensor * counts  = sc.second.second;
-
-        if (!in_sum2 || !counts) {
-            LOG_ERR("%s: mismatched sums and counts for %s\n", __func__, name.c_str());
-            gguf_free(ctx_gguf);
-            ggml_free(ctx);
-            return false;
-        }
-
-        auto & e = m_stats[name];
-
-        int64_t nval = ggml_nelements(in_sum2);
-        if (e.values.empty()) {
-            e.values.resize(nval, 0.0f);
-        } else if ((size_t) nval != e.values.size()) {
-            LOG_ERR("%s: mismatched sums size for %s: %zu != %zu\n", __func__, name.c_str(), (size_t) nval, e.values.size());
-            gguf_free(ctx_gguf);
-            ggml_free(ctx);
-            return false;
-        }
+            // GGUF format: raw sums and counts, accumulate directly
+            const int64_t nval    = entry.sums.size();
+            const int64_t ncounts = entry.counts.size();
+
+            if (e.values.empty()) {
+                e.values.resize(nval, 0.0f);
+            } else if ((size_t) nval != e.values.size()) {
+                LOG_ERR("%s: mismatched sums size for %s: %zu != %zu\n", __func__, name.c_str(), (size_t) nval, e.values.size());
+                return false;
+            }
 
-        int64_t ncounts = ggml_nelements(counts);
-        if (e.counts.empty()) {
-            e.counts.resize(ncounts, 0);
-        } else if (e.counts.size() == 1 && ncounts > 1) {
-            // broadcast, when loading an old imatrix
-            e.counts.resize(ncounts, e.counts[0]);
-        } else if ((size_t) ncounts != e.counts.size()) {
-            LOG_ERR("%s: mismatched counts size for %s: %zu != %zu\n", __func__, name.c_str(), (size_t) ncounts, e.counts.size());
-            gguf_free(ctx_gguf);
-            ggml_free(ctx);
-            return false;
-        }
+            if (e.counts.empty()) {
+                e.counts.resize(ncounts, 0);
+            } else if (e.counts.size() == 1 && ncounts > 1) {
+                e.counts.resize(ncounts, e.counts[0]);
+            } else if ((size_t) ncounts != e.counts.size()) {
+                LOG_ERR("%s: mismatched counts size for %s: %zu != %zu\n", __func__, name.c_str(), (size_t) ncounts, e.counts.size());
+                return false;
+            }
 
-        // Recreate the state as expected by save_imatrix()
-        for (int64_t j = 0; j < nval; j++) {
-            e.values[j] += ((const float *) in_sum2->data)[j];
-        }
-        for (int64_t j = 0; j < ncounts; j++) {
-            e.counts[j] += std::lround(((const float *) counts->data)[j]);
+            for (int64_t j = 0; j < nval; ++j) {
+                e.values[j] += entry.sums[j];
+            }
+            for (int64_t j = 0; j < ncounts; ++j) {
+                e.counts[j] += entry.counts[j];
+            }
         }
     }
 
-    // TODO: extract into its own method; this is also used by the legacy format
+    m_datasets.insert(m_datasets.end(), loaded.datasets.begin(), loaded.datasets.end());
+
     // Calculate the last chunk count
     int64_t max_count = 0;
     for (const auto & stats : m_stats) {
@@ -831,10 +686,8 @@ bool IMatrixCollector::load_imatrix(const char * file_name) {
             }
         }
     }
-    m_last_chunk = max_count / (m_params.n_ctx / m_params.n_parallel);
+    m_last_chunk = max_count / chunk_size;
 
-    gguf_free(ctx_gguf);
-    ggml_free(ctx);
     return true;
 }
 
@@ -1218,6 +1071,9 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    // set_params before show_statistics so load_imatrix has valid n_ctx/n_parallel
+    g_collector.set_params(params);
+
     if (params.show_statistics) {
         if (!show_statistics(params)) {
             return 1;