]> git.djapps.eu Git - pkg/ggml/sources/llama.cpp/commitdiff
tests : move save-load-state from examples to tests (#23336)
authorGeorgi Gerganov <redacted>
Thu, 21 May 2026 11:41:50 +0000 (14:41 +0300)
committerGitHub <redacted>
Thu, 21 May 2026 11:41:50 +0000 (14:41 +0300)
* tests : move save-load-state from examples to tests

- Move examples/save-load-state/ to tests/test-save-load-state.cpp
- Remove subdirectory reference from examples/CMakeLists.txt
- Add test to tests/CMakeLists.txt as a model test
- Remove CODEOWNERS entry for removed example directory

Assisted-by: llama.cpp:local pi
* cont : update ci

CODEOWNERS
ci/run.sh
examples/CMakeLists.txt
examples/save-load-state/CMakeLists.txt [deleted file]
examples/save-load-state/save-load-state.cpp [deleted file]
tests/CMakeLists.txt
tests/test-save-load-state.cpp [new file with mode: 0644]

index f92fe98ce8c1d745591f956aa7e27e69087afcc7..4b9d90177154481c41163085b16aec7e5bf5c33c 100644 (file)
@@ -49,7 +49,6 @@
 /examples/parallel/                     @ggerganov
 /examples/passkey/                      @ggerganov
 /examples/retrieval/                    @ggerganov
-/examples/save-load-state/              @ggerganov
 /examples/speculative-simple/           @ggerganov
 /examples/speculative/                  @ggerganov
 /ggml/cmake/                            @ggerganov
index a8cbd3371d3765c2a388d252e6b2c297e78944cb..b096dc23b66fcc5e1084a4afe87d9df3df1be5d2 100755 (executable)
--- a/ci/run.sh
+++ b/ci/run.sh
@@ -461,10 +461,10 @@ function gg_run_qwen3_0_6b {
 
     (time ./bin/llama-imatrix --model ${model_f16} -f ${wiki_test} -ngl 99 -c 1024 -b 512 --chunks 2 ) 2>&1 | tee -a $OUT/${ci}-imatrix.log
 
-    (time ./bin/llama-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
-    (time ./bin/llama-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
-    (time ./bin/llama-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
-    (time ./bin/llama-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
+    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa off --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
+    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 10 -c 1024 -fa on  --no-op-offload) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
+    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa off                ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
+    (time ./bin/test-save-load-state --model ${model_q4_0} -ngl 99 -c 1024 -fa on                 ) 2>&1 | tee -a $OUT/${ci}-save-load-state.log
 
     function check_ppl {
         qnt="$1"
index a29dc707c3dc82d6d0207fe481235167dda7a489..39f802d250e141135aafe05f532192e435ce135f 100644 (file)
@@ -27,7 +27,6 @@ else()
     add_subdirectory(parallel)
     add_subdirectory(passkey)
     add_subdirectory(retrieval)
-    add_subdirectory(save-load-state)
     add_subdirectory(simple)
     add_subdirectory(simple-chat)
     add_subdirectory(speculative)
diff --git a/examples/save-load-state/CMakeLists.txt b/examples/save-load-state/CMakeLists.txt
deleted file mode 100644 (file)
index 7802467..0000000
+++ /dev/null
@@ -1,5 +0,0 @@
-set(TARGET llama-save-load-state)
-add_executable(${TARGET} save-load-state.cpp)
-install(TARGETS ${TARGET} RUNTIME)
-target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT})
-target_compile_features(${TARGET} PRIVATE cxx_std_17)
diff --git a/examples/save-load-state/save-load-state.cpp b/examples/save-load-state/save-load-state.cpp
deleted file mode 100644 (file)
index 97ab7c6..0000000
+++ /dev/null
@@ -1,345 +0,0 @@
-#include "arg.h"
-#include "common.h"
-#include "log.h"
-#include "llama-cpp.h"
-
-#include <clocale>
-#include <vector>
-
-struct llama_batch_ptr {
-    llama_batch batch;
-
-    llama_batch_ptr(int32_t n_tokens, int32_t embd, int32_t n_seq_max)
-        : batch{llama_batch_init(n_tokens, embd, n_seq_max)} {}
-
-    ~llama_batch_ptr() { llama_batch_free(batch); }
-
-    llama_batch_ptr(const llama_batch_ptr &) = delete;
-    llama_batch_ptr & operator=(const llama_batch_ptr &) = delete;
-    llama_batch_ptr(llama_batch_ptr &&) = default;
-    llama_batch_ptr & operator=(llama_batch_ptr &&) = default;
-
-    llama_batch & get() { return batch; }
-    const llama_batch & get() const { return batch; }
-};
-
-static std::string generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
-    std::string result;
-    llama_batch_ptr batch(1, 0, 1);
-
-    for (int i = 0; i < n_predict; i++) {
-        auto next_token     = llama_sampler_sample(smpl, ctx, -1);
-        auto next_token_str = common_token_to_piece(ctx, next_token);
-
-        LOG("%s", next_token_str.c_str());
-        result += next_token_str;
-
-        common_batch_clear(batch.get());
-        common_batch_add(batch.get(), next_token, n_past, {seq_id}, true);
-
-        if (llama_decode(ctx, batch.get())) {
-            LOG_ERR("\n%s: failed to evaluate\n", __func__);
-            return {};
-        }
-        n_past++;
-    }
-
-    return result;
-}
-
-// Test 1: baseline
-// - tokenize the prompt
-// - decode all but the last token
-// - save state to disk
-// - decode the last token
-// - generate n_predict tokens
-static std::string test_baseline(struct llama_model * model, const struct common_params & params) {
-    auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
-
-    auto sparams = llama_sampler_chain_default_params();
-    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
-    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
-    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
-
-    auto n_past = 0;
-    if (!common_prompt_batch_decode(ctx.get(), tokens, n_past, params.n_batch, params.out_file, true)) {
-        LOG_ERR("%s: failed to decode prompt\n", __func__);
-        return {};
-    }
-
-    LOG("\n=== Test 1: baseline ===\n");
-    LOG("%s", params.prompt.c_str());
-
-    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
-    if (result.empty()) {
-        return {};
-    }
-
-    LOG("\n");
-
-    return result;
-}
-
-
-// Test 2: state load
-// - create a new context
-// - load state from file
-// - replay the last prompt token
-// - generate n_predict tokens and compare against expected result
-static bool test_state_load(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
-    auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
-
-    auto sparams = llama_sampler_chain_default_params();
-    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
-    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
-    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
-
-    LOG("\n=== Test 2: state load ===\n");
-    LOG("%s", params.prompt.c_str());
-
-    // Load state from file
-    std::vector<llama_token> unused_sts(tokens.size());
-    size_t n_token_count_out = 0;
-
-    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
-        LOG_ERR("\n%s: failed to load state\n", __func__);
-        return false;
-    }
-
-    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
-
-    // Replay last token
-    int n_past = (int) n_token_count_out;
-    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
-        return false;
-    }
-    n_past++;
-
-    // Generate tokens
-    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
-    if (result.empty()) {
-        return false;
-    }
-
-    if (result != expected_result) {
-        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
-        return false;
-    }
-
-    LOG("\nPASS\n");
-    return true;
-}
-
-
-// Test 3: seq copy (host)
-// - create a multi-seq context
-// - load state from file
-// - replay the last prompt token
-// - migrate KV cache from seq 0 to seq 1 via the CPU path
-// - generate n_predict tokens on seq 1 and compare against expected result
-static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
-    auto params_ctx = common_context_params_to_llama(params);
-    params_ctx.n_seq_max = 2;
-    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
-    auto sparams = llama_sampler_chain_default_params();
-    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
-    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
-    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
-
-    LOG("\n=== Test 3: seq copy (host) ===\n");
-    LOG("%s", params.prompt.c_str());
-
-    // Load state from file
-    std::vector<llama_token> unused_sts(tokens.size());
-    size_t n_token_count_out = 0;
-
-    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
-        LOG_ERR("\n%s: failed to load state\n", __func__);
-        return false;
-    }
-
-    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
-
-    // Replay last token
-    int n_past = (int) n_token_count_out;
-    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
-        return false;
-    }
-    n_past++;
-
-    // Migrate KV cache from seq 0 to seq 1 (CPU path)
-    {
-        std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx.get(), 0));
-        const size_t ncopy = llama_state_seq_get_data(ctx.get(), seq_store.data(), seq_store.size(), 0);
-        if (ncopy != seq_store.size()) {
-            LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
-            return false;
-        }
-        LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
-
-        llama_memory_clear(llama_get_memory(ctx.get()), true);
-        LOG_TRC("%s: kv cache cleared\n", __func__);
-
-        const size_t nset = llama_state_seq_set_data(ctx.get(), seq_store.data(), seq_store.size(), 1);
-        if (nset != seq_store.size()) {
-            LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
-            return false;
-        }
-        LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
-    }
-
-    // Generate tokens on seq 1
-    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
-    if (result.empty()) {
-        return false;
-    }
-
-    if (result != expected_result) {
-        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
-        return false;
-    }
-
-    LOG("\nPASS\n");
-    return true;
-}
-
-
-// Test 4: seq copy (device)
-// - create a multi-seq context
-// - load state from file
-// - replay the last prompt token
-// - migrate KV cache from seq 0 to seq 1 via the on-device path
-// - generate n_predict tokens on seq 1 and compare against expected result
-static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
-    auto params_ctx = common_context_params_to_llama(params);
-    params_ctx.n_seq_max = 2;
-    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
-
-    auto sparams = llama_sampler_chain_default_params();
-    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
-    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
-
-    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
-
-    LOG("\n=== Test 4: seq copy (device) ===\n");
-    LOG("%s", params.prompt.c_str());
-
-    // Load state from file
-    std::vector<llama_token> unused_sts(tokens.size());
-    size_t n_token_count_out = 0;
-
-    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
-        LOG_ERR("\n%s: failed to load state\n", __func__);
-        return false;
-    }
-
-    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
-
-    // Replay last token
-    int n_past = (int) n_token_count_out;
-    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
-        return false;
-    }
-    n_past++;
-
-    // Migrate KV cache from seq 0 to seq 1 (on-device path)
-    {
-        std::vector<uint8_t> seq_store(llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE));
-        const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
-        if (ncopy != seq_store.size()) {
-            LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
-            return false;
-        }
-        LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
-
-        llama_memory_clear(llama_get_memory(ctx.get()), true);
-        LOG_TRC("%s: kv cache cleared\n", __func__);
-
-        const size_t nset = llama_state_seq_set_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
-        if (nset != seq_store.size()) {
-            LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
-            return false;
-        }
-        LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
-    }
-
-    // Generate tokens on seq 1
-    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
-    if (result.empty()) {
-        return false;
-    }
-
-    if (result != expected_result) {
-        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
-        return false;
-    }
-
-    LOG("\nPASS\n");
-    return true;
-}
-
-
-int main(int argc, char ** argv) {
-    std::setlocale(LC_NUMERIC, "C");
-
-    common_params params;
-    params.prompt = "The quick brown fox";
-    params.out_file = "dump_state.bin";
-    params.sampling.seed = 1234;
-
-    common_init();
-
-    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
-        return 1;
-    }
-
-    if (params.n_parallel == 1) {
-        LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
-        params.kv_unified = true;
-    }
-
-    if (params.n_predict < 0) {
-        params.n_predict = 16;
-    }
-
-    ggml_backend_load_all();
-
-    auto llama_init = common_init_from_params(params, true);
-    auto * model = llama_init->model();
-
-    if (model == nullptr) {
-        LOG_ERR("%s: failed to init\n", __func__);
-        return 1;
-    }
-
-    GGML_ASSERT(llama_init->context() == nullptr);
-
-    // Test 1: baseline (saves state to disk)
-    auto result_baseline = test_baseline(model, params);
-    if (result_baseline.empty()) {
-        return 1;
-    }
-
-    // Test 2: state load
-    if (!test_state_load(model, params, result_baseline)) {
-        return 1;
-    }
-
-    // Test 3: seq copy (host)
-    if (!test_seq_cp_host(model, params, result_baseline)) {
-        return 1;
-    }
-
-    // Test 4: seq copy (device)
-    if (!test_seq_cp_device(model, params, result_baseline)) {
-        return 1;
-    }
-
-    LOG("\nAll tests passed.\n");
-
-    return 0;
-}
index 0fdbd39c94a5797a60434e2343945d1d3b1c5de2..33ae3b303cfcbbd536d02fdadd6e2af986e527d4 100644 (file)
@@ -255,6 +255,10 @@ set_tests_properties(test-state-restore-fragmented PROPERTIES FIXTURES_REQUIRED
 llama_build_and_test(test-recurrent-state-rollback.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
 set_tests_properties(test-recurrent-state-rollback PROPERTIES FIXTURES_REQUIRED test-download-model)
 
+# Test state save/load functionality
+llama_build_and_test(test-save-load-state.cpp LABEL "model" ARGS -m "${MODEL_DEST}")
+set_tests_properties(test-save-load-state PROPERTIES FIXTURES_REQUIRED test-download-model)
+
 if (NOT GGML_BACKEND_DL)
     # these tests use the backends directly and cannot be built with dynamic loading
     llama_build_and_test(test-barrier.cpp)
diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp
new file mode 100644 (file)
index 0000000..97ab7c6
--- /dev/null
@@ -0,0 +1,345 @@
+#include "arg.h"
+#include "common.h"
+#include "log.h"
+#include "llama-cpp.h"
+
+#include <clocale>
+#include <vector>
+
+struct llama_batch_ptr {
+    llama_batch batch;
+
+    llama_batch_ptr(int32_t n_tokens, int32_t embd, int32_t n_seq_max)
+        : batch{llama_batch_init(n_tokens, embd, n_seq_max)} {}
+
+    ~llama_batch_ptr() { llama_batch_free(batch); }
+
+    llama_batch_ptr(const llama_batch_ptr &) = delete;
+    llama_batch_ptr & operator=(const llama_batch_ptr &) = delete;
+    llama_batch_ptr(llama_batch_ptr &&) = default;
+    llama_batch_ptr & operator=(llama_batch_ptr &&) = default;
+
+    llama_batch & get() { return batch; }
+    const llama_batch & get() const { return batch; }
+};
+
+static std::string generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) {
+    std::string result;
+    llama_batch_ptr batch(1, 0, 1);
+
+    for (int i = 0; i < n_predict; i++) {
+        auto next_token     = llama_sampler_sample(smpl, ctx, -1);
+        auto next_token_str = common_token_to_piece(ctx, next_token);
+
+        LOG("%s", next_token_str.c_str());
+        result += next_token_str;
+
+        common_batch_clear(batch.get());
+        common_batch_add(batch.get(), next_token, n_past, {seq_id}, true);
+
+        if (llama_decode(ctx, batch.get())) {
+            LOG_ERR("\n%s: failed to evaluate\n", __func__);
+            return {};
+        }
+        n_past++;
+    }
+
+    return result;
+}
+
+// Test 1: baseline
+// - tokenize the prompt
+// - decode all but the last token
+// - save state to disk
+// - decode the last token
+// - generate n_predict tokens
+static std::string test_baseline(struct llama_model * model, const struct common_params & params) {
+    auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
+
+    auto sparams = llama_sampler_chain_default_params();
+    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
+    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
+
+    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
+
+    auto n_past = 0;
+    if (!common_prompt_batch_decode(ctx.get(), tokens, n_past, params.n_batch, params.out_file, true)) {
+        LOG_ERR("%s: failed to decode prompt\n", __func__);
+        return {};
+    }
+
+    LOG("\n=== Test 1: baseline ===\n");
+    LOG("%s", params.prompt.c_str());
+
+    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
+    if (result.empty()) {
+        return {};
+    }
+
+    LOG("\n");
+
+    return result;
+}
+
+
+// Test 2: state load
+// - create a new context
+// - load state from file
+// - replay the last prompt token
+// - generate n_predict tokens and compare against expected result
+static bool test_state_load(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
+    auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))};
+
+    auto sparams = llama_sampler_chain_default_params();
+    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
+    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
+
+    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
+
+    LOG("\n=== Test 2: state load ===\n");
+    LOG("%s", params.prompt.c_str());
+
+    // Load state from file
+    std::vector<llama_token> unused_sts(tokens.size());
+    size_t n_token_count_out = 0;
+
+    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
+        LOG_ERR("\n%s: failed to load state\n", __func__);
+        return false;
+    }
+
+    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
+
+    // Replay last token
+    int n_past = (int) n_token_count_out;
+    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
+        return false;
+    }
+    n_past++;
+
+    // Generate tokens
+    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
+    if (result.empty()) {
+        return false;
+    }
+
+    if (result != expected_result) {
+        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
+        return false;
+    }
+
+    LOG("\nPASS\n");
+    return true;
+}
+
+
+// Test 3: seq copy (host)
+// - create a multi-seq context
+// - load state from file
+// - replay the last prompt token
+// - migrate KV cache from seq 0 to seq 1 via the CPU path
+// - generate n_predict tokens on seq 1 and compare against expected result
+static bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
+    auto params_ctx = common_context_params_to_llama(params);
+    params_ctx.n_seq_max = 2;
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
+
+    auto sparams = llama_sampler_chain_default_params();
+    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
+    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
+
+    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
+
+    LOG("\n=== Test 3: seq copy (host) ===\n");
+    LOG("%s", params.prompt.c_str());
+
+    // Load state from file
+    std::vector<llama_token> unused_sts(tokens.size());
+    size_t n_token_count_out = 0;
+
+    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
+        LOG_ERR("\n%s: failed to load state\n", __func__);
+        return false;
+    }
+
+    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
+
+    // Replay last token
+    int n_past = (int) n_token_count_out;
+    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
+        return false;
+    }
+    n_past++;
+
+    // Migrate KV cache from seq 0 to seq 1 (CPU path)
+    {
+        std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx.get(), 0));
+        const size_t ncopy = llama_state_seq_get_data(ctx.get(), seq_store.data(), seq_store.size(), 0);
+        if (ncopy != seq_store.size()) {
+            LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
+            return false;
+        }
+        LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
+
+        llama_memory_clear(llama_get_memory(ctx.get()), true);
+        LOG_TRC("%s: kv cache cleared\n", __func__);
+
+        const size_t nset = llama_state_seq_set_data(ctx.get(), seq_store.data(), seq_store.size(), 1);
+        if (nset != seq_store.size()) {
+            LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
+            return false;
+        }
+        LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
+    }
+
+    // Generate tokens on seq 1
+    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
+    if (result.empty()) {
+        return false;
+    }
+
+    if (result != expected_result) {
+        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
+        return false;
+    }
+
+    LOG("\nPASS\n");
+    return true;
+}
+
+
+// Test 4: seq copy (device)
+// - create a multi-seq context
+// - load state from file
+// - replay the last prompt token
+// - migrate KV cache from seq 0 to seq 1 via the on-device path
+// - generate n_predict tokens on seq 1 and compare against expected result
+static bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const std::string & expected_result) {
+    auto params_ctx = common_context_params_to_llama(params);
+    params_ctx.n_seq_max = 2;
+    auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)};
+
+    auto sparams = llama_sampler_chain_default_params();
+    auto smpl = llama_sampler_ptr{llama_sampler_chain_init(sparams)};
+    llama_sampler_chain_add(smpl.get(), llama_sampler_init_dist(params.sampling.seed));
+
+    auto tokens = common_tokenize(ctx.get(), params.prompt, true);
+
+    LOG("\n=== Test 4: seq copy (device) ===\n");
+    LOG("%s", params.prompt.c_str());
+
+    // Load state from file
+    std::vector<llama_token> unused_sts(tokens.size());
+    size_t n_token_count_out = 0;
+
+    if (!llama_state_load_file(ctx.get(), params.out_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
+        LOG_ERR("\n%s: failed to load state\n", __func__);
+        return false;
+    }
+
+    LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
+
+    // Replay last token
+    int n_past = (int) n_token_count_out;
+    if (!common_replay_last_token(ctx.get(), tokens.back(), n_past)) {
+        return false;
+    }
+    n_past++;
+
+    // Migrate KV cache from seq 0 to seq 1 (on-device path)
+    {
+        std::vector<uint8_t> seq_store(llama_state_seq_get_size_ext(ctx.get(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE));
+        const size_t ncopy = llama_state_seq_get_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+        if (ncopy != seq_store.size()) {
+            LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
+            return false;
+        }
+        LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
+
+        llama_memory_clear(llama_get_memory(ctx.get()), true);
+        LOG_TRC("%s: kv cache cleared\n", __func__);
+
+        const size_t nset = llama_state_seq_set_data_ext(ctx.get(), seq_store.data(), seq_store.size(), 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
+        if (nset != seq_store.size()) {
+            LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
+            return false;
+        }
+        LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
+    }
+
+    // Generate tokens on seq 1
+    auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 1);
+    if (result.empty()) {
+        return false;
+    }
+
+    if (result != expected_result) {
+        LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
+        return false;
+    }
+
+    LOG("\nPASS\n");
+    return true;
+}
+
+
+int main(int argc, char ** argv) {
+    std::setlocale(LC_NUMERIC, "C");
+
+    common_params params;
+    params.prompt = "The quick brown fox";
+    params.out_file = "dump_state.bin";
+    params.sampling.seed = 1234;
+
+    common_init();
+
+    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
+        return 1;
+    }
+
+    if (params.n_parallel == 1) {
+        LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
+        params.kv_unified = true;
+    }
+
+    if (params.n_predict < 0) {
+        params.n_predict = 16;
+    }
+
+    ggml_backend_load_all();
+
+    auto llama_init = common_init_from_params(params, true);
+    auto * model = llama_init->model();
+
+    if (model == nullptr) {
+        LOG_ERR("%s: failed to init\n", __func__);
+        return 1;
+    }
+
+    GGML_ASSERT(llama_init->context() == nullptr);
+
+    // Test 1: baseline (saves state to disk)
+    auto result_baseline = test_baseline(model, params);
+    if (result_baseline.empty()) {
+        return 1;
+    }
+
+    // Test 2: state load
+    if (!test_state_load(model, params, result_baseline)) {
+        return 1;
+    }
+
+    // Test 3: seq copy (host)
+    if (!test_seq_cp_host(model, params, result_baseline)) {
+        return 1;
+    }
+
+    // Test 4: seq copy (device)
+    if (!test_seq_cp_device(model, params, result_baseline)) {
+        return 1;
+    }
+
+    LOG("\nAll tests passed.\n");
+
+    return 0;
+}