server: genes (LoRA adapters) page between the device and host RAM (stacked on #70)

common/arg.cpp

@@ -330,6 +330,18 @@ static std::string get_all_kv_cache_types() {
     return msg.str();
 }
 
+// an exact byte count: the caller passes a measured number, so no unit suffixes are guessed
+static uint64_t parse_byte_count(const std::string & value, const char * flag) {
+    if (value.empty() || value.size() > 20 || value.find_first_not_of("0123456789") != std::string::npos) {
+        throw std::invalid_argument(std::string(flag) + " takes an unsigned byte count, got: " + value);
+    }
+    try {
+        return std::stoull(value);
+    } catch (const std::exception &) {
+        throw std::invalid_argument(std::string(flag) + " is out of range: " + value);
+    }
+}
+
 static bool parse_bool_value(const std::string & value) {
     if (is_truthy(value)) {
         return true;
@@ -1691,6 +1703,23 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.n_ctx_checkpoints = value;
         }
     ).set_env("LLAMA_ARG_CTX_CHECKPOINTS").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}));
+    add_opt(common_arg(
+        {"--ctx-checkpoints-bytes"}, "BYTES",
+        "byte budget for the context checkpoints of all slots together, in place of the per-slot count "
+        "(default: 0 = use --ctx-checkpoints); when a new checkpoint would exceed it, the slot holding the "
+        "most checkpoint bytes drops its oldest checkpoint after its first, so each slot keeps its head and its newest",
+        [](common_params & params, const std::string & value) {
+            params.ctx_checkpoints_bytes = parse_byte_count(value, "--ctx-checkpoints-bytes");
+        }
+    ).set_env("LLAMA_ARG_CTX_CHECKPOINTS_BYTES").set_examples({LLAMA_EXAMPLE_SERVER}));
+    add_opt(common_arg(
+        {"--page-ram-bytes"}, "BYTES",
+        "ceiling of the host-RAM page store that /slots page_out fills (default: 0 = no ceiling); "
+        "a page_out that would exceed it is refused, so the caller spills a page to disk first",
+        [](common_params & params, const std::string & value) {
+            params.page_ram_bytes = parse_byte_count(value, "--page-ram-bytes");
+        }
+    ).set_env("LLAMA_ARG_PAGE_RAM_BYTES").set_examples({LLAMA_EXAMPLE_SERVER}));
     add_opt(common_arg(
         {"-cms", "--checkpoint-min-step"}, "N",
         string_format("minimum spacing between context checkpoints in tokens (default: %d, 0 = no minimum)", params.checkpoint_min_step),

common/common.h

@@ -621,6 +621,8 @@ struct common_params {
     bool    cache_idle_slots    = true;  // save and clear idle slots upon starting a new task
     int32_t n_ctx_checkpoints   = 32;    // max number of context checkpoints per slot
     int32_t checkpoint_min_step = 8192;  // minimum spacing between context checkpoints
+    uint64_t ctx_checkpoints_bytes = 0;  // checkpoint byte budget across all slots; 0 = the count above governs
+    uint64_t page_ram_bytes        = 0;  // ceiling of the named host-RAM page store; 0 = no ceiling
     int32_t cache_ram_mib       = 8192;  // -1 = no limit, 0 - disable, 1 = 1 MiB, etc.
 
     std::string hostname      = "127.0.0.1";

include/llama.h

@@ -796,6 +796,20 @@ extern "C" {
     // NOTE: loaded adapters that are not manually freed will be freed when the associated model is deleted
     LLAMA_API void llama_adapter_lora_free(struct llama_adapter_lora * adapter);
 
+    // Continuum fork: page an adapter's tensors between their device buffers and host memory, so a
+    // loaded adapter no context applies costs host RAM instead of device memory, and comes back
+    // with a copy instead of a read from disk. The adapter stays loaded and keeps its identity;
+    // a paged-out adapter must be paged in before a context applies it. The caller serializes
+    // these with every use of the adapter (synchronize the contexts that ran it first).
+    // Return the bytes moved; 0 when it is already there; -1 when an allocation failed (the
+    // adapter is then left exactly as it was).
+    LLAMA_API int64_t llama_adapter_lora_page_out(struct llama_adapter_lora * adapter);
+    LLAMA_API int64_t llama_adapter_lora_page_in (struct llama_adapter_lora * adapter);
+    LLAMA_API bool    llama_adapter_lora_is_paged_out(const struct llama_adapter_lora * adapter);
+    // The bytes the adapter's tensors take where they are now: its device buffers when resident,
+    // host memory when paged out.
+    LLAMA_API uint64_t llama_adapter_lora_n_bytes(const struct llama_adapter_lora * adapter);
+
     // Get the invocation tokens if the current lora is an alora
     LLAMA_API uint64_t            llama_adapter_get_alora_n_invocation_tokens(const struct llama_adapter_lora * adapter);
     LLAMA_API const llama_token * llama_adapter_get_alora_invocation_tokens  (const struct llama_adapter_lora * adapter);

src/llama-adapter.cpp

@@ -6,6 +6,7 @@
 
 #include <map>
 #include <cassert>
+#include <functional>
 #include <mutex>
 #include <shared_mutex>
 #include <sstream>
@@ -479,6 +480,97 @@ int32_t llama_adapter_meta_val_str_by_index(const llama_adapter_lora * adapter,
     return snprintf(buf, buf_size, "%s", it->second.c_str());
 }
 
+// Continuum fork: adapter paging between device buffers and host memory (llama.h).
+
+static void lora_for_each_tensor(const llama_adapter_lora & adapter, const std::function<void(size_t, ggml_tensor *)> & fn) {
+    for (size_t i = 0; i < adapter.ctxs.size(); ++i) {
+        for (ggml_tensor * t = ggml_get_first_tensor(adapter.ctxs[i].get()); t; t = ggml_get_next_tensor(adapter.ctxs[i].get(), t)) {
+            fn(i, t);
+        }
+    }
+}
+
+int64_t llama_adapter_lora_page_out(llama_adapter_lora * adapter) {
+    if (adapter->paged) {
+        return 0;
+    }
+    auto page = std::make_unique<llama_adapter_lora_host_page>();
+    page->bufts.assign(adapter->ctxs.size(), nullptr);
+    size_t total = 0;
+    lora_for_each_tensor(*adapter, [&](size_t i, ggml_tensor * t) {
+        if (page->bufts[i] == nullptr && t->buffer != nullptr) {
+            page->bufts[i] = ggml_backend_buffer_get_type(t->buffer);
+        }
+        total += ggml_nbytes(t);
+    });
+    try {
+        page->host.resize(total);
+    } catch (const std::bad_alloc &) {
+        LLAMA_LOG_ERROR("%s: no host memory for %zu bytes of LoRA tensors\n", __func__, total);
+        return -1;
+    }
+    size_t off = 0;
+    lora_for_each_tensor(*adapter, [&](size_t, ggml_tensor * t) {
+        ggml_backend_tensor_get(t, page->host.data() + off, 0, ggml_nbytes(t));
+        off += ggml_nbytes(t);
+    });
+    adapter->bufs.clear(); // frees the device buffers
+    lora_for_each_tensor(*adapter, [](size_t, ggml_tensor * t) {
+        t->data   = nullptr;
+        t->buffer = nullptr;
+    });
+    adapter->paged = std::move(page);
+    return (int64_t) total;
+}
+
+int64_t llama_adapter_lora_page_in(llama_adapter_lora * adapter) {
+    if (!adapter->paged) {
+        return 0;
+    }
+    const llama_adapter_lora_host_page & page = *adapter->paged;
+    std::vector<ggml_backend_buffer_ptr> bufs;
+    for (size_t i = 0; i < adapter->ctxs.size(); ++i) {
+        if (page.bufts[i] == nullptr) {
+            continue; // a context with no tensors
+        }
+        ggml_backend_buffer_ptr buf { ggml_backend_alloc_ctx_tensors_from_buft(adapter->ctxs[i].get(), page.bufts[i]) };
+        if (!buf) {
+            LLAMA_LOG_ERROR("%s: failed to re-allocate a %s LoRA buffer\n", __func__, ggml_backend_buft_name(page.bufts[i]));
+            bufs.clear();
+            lora_for_each_tensor(*adapter, [](size_t, ggml_tensor * t) {
+                t->data   = nullptr;
+                t->buffer = nullptr;
+            });
+            return -1;
+        }
+        bufs.emplace_back(std::move(buf));
+    }
+    size_t off = 0;
+    lora_for_each_tensor(*adapter, [&](size_t, ggml_tensor * t) {
+        ggml_backend_tensor_set(t, page.host.data() + off, 0, ggml_nbytes(t));
+        off += ggml_nbytes(t);
+    });
+    adapter->bufs = std::move(bufs);
+    const int64_t moved = (int64_t) page.host.size();
+    adapter->paged.reset();
+    return moved;
+}
+
+bool llama_adapter_lora_is_paged_out(const llama_adapter_lora * adapter) {
+    return adapter->paged != nullptr;
+}
+
+uint64_t llama_adapter_lora_n_bytes(const llama_adapter_lora * adapter) {
+    if (adapter->paged) {
+        return adapter->paged->host.size();
+    }
+    uint64_t n = 0;
+    for (const auto & buf : adapter->bufs) {
+        n += ggml_backend_buffer_get_size(buf.get());
+    }
+    return n;
+}
+
 void llama_adapter_lora_free(llama_adapter_lora * adapter) {
     if (adapter == nullptr) {
         return;

src/llama-adapter.h

@@ -60,6 +60,15 @@ struct llama_adapter_lora_weight {
     llama_adapter_lora_weight(ggml_tensor * a, ggml_tensor * b) : a(a), b(b) {}
 };
 
+// Continuum fork: an adapter's tensors held in host memory while its device buffers are freed
+// (llama_adapter_lora_page_out). `bufts` is aligned with `ctxs`: the buffer type each context's
+// tensors are re-allocated from on page-in; `host` is every tensor's bytes, contexts in order and
+// tensors in context order.
+struct llama_adapter_lora_host_page {
+    std::vector<ggml_backend_buffer_type_t> bufts;
+    std::vector<uint8_t> host;
+};
+
 struct llama_adapter_lora {
     llama_model * model = nullptr;
 
@@ -79,6 +88,9 @@ struct llama_adapter_lora {
     // activated lora (aLoRA)
     std::vector<llama_token> alora_invocation_tokens;
 
+    // non-null while paged out to host memory: the tensors then have no device data
+    std::unique_ptr<llama_adapter_lora_host_page> paged;
+
     explicit llama_adapter_lora(llama_model * model) : model(model) {}
     ~llama_adapter_lora() = default;
 

src/llama-context.cpp

@@ -3995,6 +3995,36 @@ size_t llama_state_seq_save_file(llama_context * ctx, const char * filepath, lla
     }
 }
 
+size_t llama_state_seq_save_file_from_data(const char * filepath, const uint8_t * data, size_t size, const llama_token * tokens, size_t n_token_count) {
+    // get_data writes [io_magic][seq_id][state]; the file is [magic][version][count][tokens][state]
+    const size_t prefix = sizeof(io_magic) + sizeof(llama_seq_id);
+    if (data == nullptr || size <= prefix) {
+        LLAMA_LOG_ERROR("%s: no sequence state to write\n", __func__);
+        return 0;
+    }
+    uint32_t magic_read = 0;
+    std::memcpy(&magic_read, data, sizeof(magic_read));
+    if (magic_read != io_magic) {
+        LLAMA_LOG_ERROR("%s: the bytes are not a sequence state (magic %08x)\n", __func__, magic_read);
+        return 0;
+    }
+    try {
+        llama_file file(filepath, "wb");
+        const uint32_t magic   = LLAMA_STATE_SEQ_MAGIC;
+        const uint32_t version = LLAMA_STATE_SEQ_VERSION;
+        const uint32_t count   = (uint32_t) n_token_count;
+        file.write_raw(&magic,   sizeof(magic));
+        file.write_raw(&version, sizeof(version));
+        file.write_raw(&count,   sizeof(count));
+        file.write_raw(tokens,   sizeof(llama_token) * n_token_count);
+        file.write_raw(data + prefix, size - prefix);
+        return file.tell();
+    } catch (const std::exception & err) {
+        LLAMA_LOG_ERROR("%s: error writing sequence state file: %s\n", __func__, err.what());
+        return 0;
+    }
+}
+
 size_t llama_state_seq_load_file(llama_context * ctx, const char * filepath, llama_seq_id dest_seq_id, llama_token * tokens_out, size_t n_token_capacity, size_t * n_token_count_out) {
     return llama_state_seq_load_file_with_progress(ctx, filepath, dest_seq_id, tokens_out, n_token_capacity, n_token_count_out, nullptr, nullptr);
 }

src/llama-ext.h

@@ -130,3 +130,19 @@ LLAMA_API uint32_t        llama_model_target_layer_ids_n(const struct llama_mode
 // if out is nullptr, returns the number of tokens without writing to out
 // caller must allocate enough memory for out before calling
 LLAMA_API uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out);
+
+//
+// sequence state pages (Continuum fork)
+//
+
+// Writes a sequence-state file (the llama_state_seq_save_file format) from bytes that
+// llama_state_seq_get_data_ext returned with LLAMA_STATE_SEQ_FLAGS_NONE, without a context:
+// a page held in host RAM spills to disk off the engine thread, and llama_state_seq_load_file
+// reads it back. Returns the bytes written, or 0 when the bytes are not a sequence state or
+// the write fails.
+LLAMA_API size_t llama_state_seq_save_file_from_data(
+        const char        * filepath,
+        const uint8_t     * data,
+        size_t              size,
+        const llama_token * tokens,
+        size_t              n_token_count);

tests/CMakeLists.txt

@@ -172,6 +172,10 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS)
     llama_build_and_test(test-chat.cpp WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})
     target_include_directories(test-chat PRIVATE ${PROJECT_SOURCE_DIR}/tools/server)
     target_link_libraries(test-chat PRIVATE server-context)
+    # Continuum fork: the RAM page store's ceiling and the page files' whole-generation publication
+    llama_build_and_test(test-server-pages.cpp)
+    target_include_directories(test-server-pages PRIVATE ${PROJECT_SOURCE_DIR}/tools/server)
+    target_link_libraries(test-server-pages PRIVATE server-context)
     # TODO: disabled on loongarch64 because the ggml-ci node lacks Python 3.8
     if (NOT ${CMAKE_SYSTEM_PROCESSOR} MATCHES "loongarch64")
         llama_build_and_test(test-json-schema-to-grammar.cpp   WORKING_DIRECTORY ${PROJECT_SOURCE_DIR})