[mbuffer] Expose single-buffer loading to Llama interface

Add new C++ function to Llama main header to load from a single memory buffer, and propagate changes to internal calls/constructors.

Author: jesusmb1995

include/llama-cpp.h

@@ -5,6 +5,7 @@
 #endif
 
 #include <memory>
+#include <vector>
 
 #include "llama.h"
 
@@ -28,3 +29,5 @@ typedef std::unique_ptr<llama_model, llama_model_deleter> llama_model_ptr;
 typedef std::unique_ptr<llama_context, llama_context_deleter> llama_context_ptr;
 typedef std::unique_ptr<llama_sampler, llama_sampler_deleter> llama_sampler_ptr;
 typedef std::unique_ptr<llama_adapter_lora, llama_adapter_lora_deleter> llama_adapter_lora_ptr;
+
+struct llama_model * llama_model_load_from_buffer(std::vector<uint8_t> && data, struct llama_model_params params);

src/llama-model-loader.cpp

@@ -8,6 +8,7 @@
 #include <cinttypes>
 #include <cstring>
 #include <future>
+#include <stdexcept>
 
 static const size_t kiB = 1024;
 static const size_t MiB = 1024*kiB;
@@ -493,8 +494,7 @@ namespace GGUFMeta {
     }
 
 llama_model_loader::llama_model_loader(
-        const std::string & fname,
-        std::vector<std::string> & splits,
+        load_input_t load_input,
         bool use_mmap,
         bool check_tensors,
         const llama_model_kv_override * param_overrides_p,
@@ -513,7 +513,7 @@ llama_model_loader::llama_model_loader(
     tensor_buft_overrides = param_tensor_buft_overrides_p;
 
     struct ggml_context * ctx = NULL;
-    gguf_file_load main_gguf(&ctx, load_input_variant::fname_load_input{fname, splits});
+    gguf_file_load main_gguf(&ctx, load_input);
     process_loaded_gguf(ctx, main_gguf, 0);
 
     meta = std::move(main_gguf.meta);
@@ -525,18 +525,21 @@ llama_model_loader::llama_model_loader(
     get_key(llm_kv(LLM_KV_SPLIT_COUNT), n_split, false);
 
     // Load additional GGML contexts
-    if (n_split > 1) {
+    if (load_input_variant::variant_supports_split_load(load_input) && n_split > 1) {
+        load_input_variant::fname_load_input base_split = load_input_variant::split_name_from_variant(load_input);
+        std::vector<std::string> &           splits     = base_split.splits;
+
         // make sure the main file is loaded first
         uint16_t idx = 0;
         const std::string kv_split_no = llm_kv(LLM_KV_SPLIT_NO);
         get_key(kv_split_no, idx);
         if (idx != 0) {
-            throw std::runtime_error(format("illegal split file idx: %d (file: %s), model must be loaded with the first split", idx, fname.c_str()));
+            throw std::runtime_error(format("illegal split file idx: %d (file: %s), model must be loaded with the first split", idx, base_split.fname.c_str()));
         }
 
         // generate list of splits if needed
         if (splits.empty()) {
-            splits = llama_get_list_splits(fname, idx, n_split);
+            splits = llama_get_list_splits(base_split.fname, idx, n_split);
         }
 
         // in case user give a custom list of splits, check if it matches the expected number
@@ -589,7 +592,7 @@ llama_model_loader::llama_model_loader(
     fver = (enum llama_fver) gguf_get_version(meta.get());
 
     LLAMA_LOG_INFO("%s: loaded meta data with %d key-value pairs and %d tensors from %s (version %s)\n",
-            __func__, n_kv, n_tensors, fname.c_str(), llama_file_version_name(fver));
+            __func__, n_kv, n_tensors, load_input_variant::identifier(load_input), llama_file_version_name(fver));
 
     // determine file type based on the number of tensors for each quantization and print meta data
     // TODO: make optional

src/llama-model-loader.h

@@ -95,8 +95,7 @@ struct llama_model_loader {
     void process_loaded_gguf(struct ggml_context * ctx, gguf_file_load & gguf_load, uint16_t idx);
 
     llama_model_loader(
-        const std::string & fname,
-        std::vector<std::string> & splits, // optional, only need if the split does not follow naming scheme
+        load_input_t load_input,
         bool use_mmap,
         bool check_tensors,
         const llama_model_kv_override * param_overrides_p,

src/llama-quant.cpp

@@ -583,7 +583,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std::
     }
 
     std::vector<std::string> splits = {};
-    llama_model_loader ml(fname_inp, splits, use_mmap, /*check_tensors*/ true, kv_overrides, nullptr);
+    load_input_variant::fname_load_input inp{fname_inp, splits};
+    llama_model_loader                   ml(inp, use_mmap, /*check_tensors*/ true, kv_overrides, nullptr);
     ml.init_mappings(false); // no prefetching
 
     llama_model model(llama_model_default_params());

src/llama.cpp

@@ -9,18 +9,24 @@
 
 #include "ggml.h"
 #include "ggml-backend.h"
+#include "uint8-buff-stream.h"
 
 #include <algorithm>
 #include <cstddef>
 #include <cstdint>
 #include <cstdio>
 #include <cstring>
 #include <ctime>
+#include <stdexcept>
 
 #if defined(_MSC_VER)
 #pragma warning(disable: 4244 4267) // possible loss of data
 #endif
 
+#ifdef __cplusplus
+#include "llama-cpp.h"
+#endif
+
 //
 // interface implementation
 //
@@ -84,7 +90,7 @@ int64_t llama_time_us(void) {
 }
 
 // Returns 0 on success, -1 on error, and -2 on cancellation via llama_progress_callback
-static int llama_model_load(const std::string & fname, std::vector<std::string> & splits, llama_model & model, llama_model_params & params) {
+static int llama_model_load(llama_model_loader & ml, llama_model & model, llama_model_params & params) {
     // loading time will be recalculated after the first eval, so
     // we take page faults deferred by mmap() into consideration
     model.t_load_us = 0;
@@ -93,8 +99,6 @@ static int llama_model_load(const std::string & fname, std::vector<std::string>
     model.t_start_us = tm.t_start_us;
 
     try {
-        llama_model_loader ml(fname, splits, params.use_mmap, params.check_tensors, params.kv_overrides, params.tensor_buft_overrides);
-
         ml.print_info();
 
         model.hparams.vocab_only = params.vocab_only;
@@ -135,8 +139,7 @@ static int llama_model_load(const std::string & fname, std::vector<std::string>
 }
 
 static struct llama_model * llama_model_load_from_file_impl(
-        const std::string & path_model,
-        std::vector<std::string> & splits,
+        llama_model_loader& ml,
         struct llama_model_params params) {
     ggml_time_init();
 
@@ -218,7 +221,7 @@ static struct llama_model * llama_model_load_from_file_impl(
         LLAMA_LOG_INFO("%s: using device %s (%s) - %zu MiB free\n", __func__, ggml_backend_dev_name(dev), ggml_backend_dev_description(dev), free/1024/1024);
     }
 
-    const int status = llama_model_load(path_model, splits, *model, params);
+    const int status = llama_model_load(ml, *model, params);
     GGML_ASSERT(status <= 0);
     if (status < 0) {
         if (status == -1) {
@@ -241,11 +244,34 @@ struct llama_model * llama_load_model_from_file(
     return llama_model_load_from_file(path_model, params);
 }
 
-struct llama_model * llama_model_load_from_file(
-        const char * path_model,
-        struct llama_model_params params) {
+static llama_model_loader create_disk_fileloader(const char * path_model, std::vector<std::string> & splits,
+                                                 struct llama_model_params params) {
+    load_input_variant::fname_load_input loader_input{ path_model, splits };
+    return llama_model_loader(loader_input, params.use_mmap, params.check_tensors, params.kv_overrides,
+                              params.tensor_buft_overrides);
+}
+
+struct llama_model * llama_model_load_from_file(const char * path_model, struct llama_model_params params) {
     std::vector<std::string> splits = {};
-    return llama_model_load_from_file_impl(path_model, splits, params);
+    llama_model_loader       ml     = create_disk_fileloader(path_model, splits, params);
+    return llama_model_load_from_file_impl(ml, params);
+}
+
+namespace {
+void override_and_disable_mmap(struct llama_model_params & params) {
+    if (params.use_mmap) {
+        LLAMA_LOG_WARN("Overriding and disabling memory mapping when loading from memory buffer\n");
+        params.use_mmap = false;
+    }
+}
+}  // namespace
+
+struct llama_model * llama_model_load_from_buffer(std::vector<uint8_t> && data, struct llama_model_params params) {
+    std::unique_ptr<std::basic_streambuf<uint8_t>> streambuf = std::make_unique<Uint8BufferStreamBuf>(std::move(data));
+    override_and_disable_mmap(params);
+    llama_model_loader ml(load_input_variant::buffer_load_input{ streambuf }, params.use_mmap, params.check_tensors,
+                          params.kv_overrides, params.tensor_buft_overrides);
+    return llama_model_load_from_file_impl(ml, params);
 }
 
 namespace {
@@ -268,7 +294,8 @@ struct llama_model * llama_model_load_from_splits(const char ** paths, size_t n_
     if (splits.empty()) {
         return nullptr;
     }
-    return llama_model_load_from_file_impl(splits.front(), splits, params);
+    llama_model_loader ml = create_disk_fileloader(splits.front().c_str(), splits, params);
+    return llama_model_load_from_file_impl(ml, params);
 }
 
 void llama_model_save_to_file(const struct llama_model * model, const char * path_model) {