add functions to get and set the whole llama state (rng, logits, embedding and kv_cache)

Author: xaedes

llama.cpp

@@ -16,7 +16,8 @@
 #include <climits>
 #include <memory>
 #include <algorithm>
-#include <initializer_list>
+#include <sstream>
+
 
 #define LLAMA_USE_SCRATCH
 #define LLAMA_MAX_SCRATCH_BUFFERS 16
@@ -1931,3 +1932,126 @@ const char * llama_print_system_info(void) {
 std::vector<std::pair<std::string, struct ggml_tensor *>>& llama_internal_get_tensor_map(struct llama_context * ctx) {
     return ctx->model.tensors_by_name;
 }
+
+// Returns the size of the state
+size_t llama_get_state_size(struct llama_context * ctx) {
+    const size_t s_bool = sizeof(int32_t);
+    // we don't know size of rng until we actually serialize it. so reserve more than enough memory for its serialized state. 
+    // for reference, std::mt19937(1337) serializes to 6701 bytes.
+    const size_t s_rng_size = sizeof(size_t); 
+    const size_t s_rng = 64*1024; 
+    const size_t s_logits_capacity = sizeof(size_t);
+    const size_t s_logits_size = sizeof(size_t);
+    const size_t s_logits = ctx->logits.capacity() * sizeof(float);
+    const size_t s_embedding_size = sizeof(size_t);
+    const size_t s_embedding = ctx->embedding.size() * sizeof(float);
+    const size_t s_kv_size = sizeof(size_t);
+    const size_t s_kv_ntok = sizeof(int);
+    const size_t s_kv = llama_get_kv_cache_size(ctx);
+    const size_t s_total = (
+        + s_rng_size
+        + s_rng
+        + s_logits_capacity
+        + s_logits_size
+        + s_logits
+        + s_embedding_size
+        + s_embedding
+        + s_kv_size
+        + s_kv_ntok
+        + s_kv
+    );
+    return s_total;
+}
+
+// Copies the state to the specified destination address
+size_t llama_copy_state_data(struct llama_context * ctx, uint8_t * dest) {
+    std::stringstream rng_ss;
+    rng_ss << ctx->rng;
+    const size_t rng_size = rng_ss.str().size();
+    char rng_buf[64*1024];
+    memset(&rng_buf[0], 0, 64*1024);
+    memcpy(&rng_buf[0], rng_ss.str().data(), rng_ss.str().size());
+    const int32_t has_evaluated_once = ctx->has_evaluated_once ? 1 : 0;
+    const int32_t logits_all = ctx->logits_all ? 1 : 0;
+    const size_t logits_capacity = ctx->logits.capacity();
+    const size_t logits_size = ctx->logits.size();
+    const size_t embedding_size = ctx->embedding.size();
+    const size_t kv_size = llama_get_kv_cache_size(ctx);
+    const int kv_ntok = llama_get_kv_cache_token_count(ctx);
+
+    uint8_t * out = dest;
+    memcpy(out, &rng_size, sizeof(size_t)); out += sizeof(size_t);
+    memcpy(out, &rng_buf[0], 64*1024); out += 64*1024;
+    memcpy(out, &logits_capacity, sizeof(size_t)); out += sizeof(size_t);
+    memcpy(out, &logits_size, sizeof(size_t)); out += sizeof(size_t);
+    if (logits_size) {
+        memcpy(out, ctx->logits.data(), logits_size * sizeof(float)); 
+    }
+    out += logits_capacity * sizeof(float);
+    memcpy(out, &embedding_size, sizeof(size_t)); out += sizeof(size_t);
+    if (embedding_size) {
+        memcpy(out, ctx->embedding.data(), embedding_size * sizeof(float)); out += embedding_size * sizeof(float);
+    }
+    memcpy(out, &kv_size, sizeof(size_t)); out += sizeof(size_t);
+    memcpy(out, &kv_ntok, sizeof(int)); out += sizeof(int);
+    if (kv_size) {
+        memcpy(out, llama_get_kv_cache(ctx), kv_size); out += kv_size;
+    }
+    const size_t written = out - dest;
+    const size_t expected = llama_get_state_size(ctx);
+    LLAMA_ASSERT(written == expected);
+    return written;
+}
+
+// Copies the state to the specified destination address
+size_t llama_set_state_data(struct llama_context * ctx, const uint8_t * src) {
+    size_t rng_size;
+    char rng_buf[64*1024];
+    std::stringstream rng_ss;
+
+    const uint8_t * in = src;
+    memcpy(&rng_size, in, sizeof(size_t)); in += sizeof(size_t);
+    memcpy(&rng_buf[0], in, 64*1024); in += 64*1024;
+    rng_ss.str(std::string(&rng_buf[0], rng_size));
+    rng_ss >> ctx->rng;
+    LLAMA_ASSERT(rng_ss.fail() == false);
+
+    int32_t has_evaluated_once;
+    int32_t logits_all;
+    size_t logits_capacity;
+    size_t logits_size;
+    size_t embedding_size;
+    size_t kv_size;
+    int kv_ntok;
+
+    memcpy(&logits_capacity, in, sizeof(size_t)); in += sizeof(size_t);
+    memcpy(&logits_size, in, sizeof(size_t)); in += sizeof(size_t);
+    LLAMA_ASSERT(ctx->logits.capacity() == logits_capacity);
+    if (logits_size) {
+        ctx->logits.resize(logits_size);
+        memcpy(ctx->logits.data(), in, logits_size * sizeof(float)); 
+    }
+    in += logits_capacity * sizeof(float);
+    memcpy(&embedding_size, in, sizeof(size_t)); in += sizeof(size_t);
+    LLAMA_ASSERT(ctx->embedding.capacity() == embedding_size);
+    if (embedding_size) {
+        memcpy(ctx->embedding.data(), in, embedding_size * sizeof(float)); 
+        in += embedding_size * sizeof(float);
+    }
+    memcpy(&kv_size, in, sizeof(size_t)); in += sizeof(size_t);
+    memcpy(&kv_ntok, in, sizeof(int)); in += sizeof(int);
+    if (kv_size) {
+        LLAMA_ASSERT(ctx->model.kv_self.buf.size == kv_size);
+        void * k_data = ctx->model.kv_self.k->data; // remember data pointers
+        void * v_data = ctx->model.kv_self.v->data; // because their value is stored in buf and overwritten by memcpy
+        memcpy(ctx->model.kv_self.buf.addr, in, kv_size); 
+        ctx->model.kv_self.k->data = k_data; // restore correct data pointers
+        ctx->model.kv_self.v->data = v_data;
+        in += kv_size;
+    }
+    ctx->model.kv_self.n = kv_ntok;
+    const size_t nread = in - src;
+    const size_t expected = llama_get_state_size(ctx);
+    LLAMA_ASSERT(nread == expected);
+    return nread;
+}

llama.h

@@ -104,6 +104,17 @@ extern "C" {
                           size_t   n_size,
                              int   n_token_count);
 
+    // Returns the size of the state
+    LLAMA_API size_t llama_get_state_size(struct llama_context * ctx);
+
+    // Copies the state to the specified destination address
+    // Returns the number of bytes copied
+    LLAMA_API size_t llama_copy_state_data(struct llama_context * ctx, uint8_t * dest);
+
+    // Copies the state to the specified destination address
+    // Returns the number of bytes read
+    LLAMA_API size_t llama_set_state_data(struct llama_context * ctx, const uint8_t * src);
+
     // Run the llama inference to obtain the logits and probabilities for the next token.
     // tokens + n_tokens is the provided batch of new tokens to process
     // n_past is the number of tokens to use from previous eval calls

py/llama_cpp/llama.py

@@ -105,6 +105,15 @@ class llama_context_params(Structure):
 lib.llama_set_kv_cache.argtypes = [llama_context_p, c_ubyte_p, c_size_t, c_int]
 lib.llama_set_kv_cache.restype = None
 
+lib.llama_get_state_size.argtypes = [llama_context_p]
+lib.llama_get_state_size.restype = c_size_t
+
+lib.llama_copy_state_data.argtypes = [llama_context_p, c_ubyte_p]
+lib.llama_copy_state_data.restype = c_size_t
+
+lib.llama_set_state_data.argtypes = [llama_context_p, c_ubyte_p]
+lib.llama_set_state_data.restype = c_size_t
+
 # Python functions
 def llama_context_default_params() -> llama_context_params:
     params = lib.llama_context_default_params()
@@ -186,3 +195,12 @@ def llama_get_kv_cache_token_count(ctx: llama_context_p) -> c_int:
 
 def llama_set_kv_cache(ctx: llama_context_p, data: c_ubyte_p, n_size:c_size_t, n_token_count:c_int):
     return lib.llama_set_kv_cache(ctx, data, n_size, n_token_count)
+
+def llama_get_state_size(ctx: llama_context_p) -> c_size_t:
+    return lib.llama_get_state_size(ctx)
+
+def llama_copy_state_data(ctx: llama_context_p, dst: c_ubyte_p) -> c_size_t:
+    return lib.llama_copy_state_data(ctx, dst)
+
+def llama_set_state_data(ctx: llama_context_p, src: c_ubyte_p) -> c_size_t:
+    return lib.llama_set_state_data(ctx, src)