#include "arg.h"
#include "common.h"
-#include "llama.h"
+#include "log.h"
+#include "llama-cpp.h"
#include <clocale>
#include <vector>
-#include <cstdio>
+struct llama_batch_ptr {
+ llama_batch batch;
-int main(int argc, char ** argv) {
- std::setlocale(LC_NUMERIC, "C");
+ 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)} {}
- common_params params;
+ ~llama_batch_ptr() { llama_batch_free(batch); }
- params.prompt = "The quick brown fox";
- params.sampling.seed = 1234;
+ 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;
- const std::string_view state_file = "dump_state.bin";
+ llama_batch & get() { return batch; }
+ const llama_batch & get() const { return batch; }
+};
- common_init();
+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);
- if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
- return 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);
- if (params.n_parallel == 1) {
- // the example uses 2 sequences, so when n_parallel == 1, we need to enable unified kv cache
- printf("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
- params.kv_unified = true;
- }
+ LOG("%s", next_token_str.c_str());
+ result += next_token_str;
- if (params.n_predict < 0) {
- params.n_predict = 16;
+ 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++;
}
- auto n_past = 0;
+ return result;
+}
- std::string result0;
- std::string result1;
- std::string result2;
- std::string result3;
+// 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))};
- // init
+ 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));
- ggml_backend_load_all();
+ auto tokens = common_tokenize(ctx.get(), params.prompt, true);
- auto llama_init = common_init_from_params(params);
+ 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 {};
+ }
- auto * model = llama_init->model();
- auto * ctx = llama_init->context();
+ LOG("\n=== Test 1: baseline ===\n");
+ LOG("%s", params.prompt.c_str());
- if (model == nullptr || ctx == nullptr) {
- fprintf(stderr, "%s : failed to init\n", __func__);
- return 1;
+ auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0);
+ if (result.empty()) {
+ return {};
}
- auto sparams = llama_sampler_chain_default_params();
+ LOG("\n");
- llama_sampler * smpl = llama_sampler_chain_init(sparams);
+ return result;
+}
- llama_sampler_chain_add(smpl, llama_sampler_init_dist(params.sampling.seed));
- // tokenize prompt
- auto tokens = common_tokenize(ctx, params.prompt, true);
+// 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))};
- const bool save_state = true;
- if (!common_prompt_batch_decode(ctx, tokens, n_past, params.n_batch, state_file, save_state)) {
- return 1;
- }
+ 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));
- // first run
- printf("\nfirst run: %s", params.prompt.c_str());
+ auto tokens = common_tokenize(ctx.get(), params.prompt, true);
- llama_batch batch = llama_batch_init(1, 0, 1);
+ LOG("\n=== Test 2: state load ===\n");
+ LOG("%s", params.prompt.c_str());
- for (auto i = 0; i < params.n_predict; i++) {
- auto next_token = llama_sampler_sample(smpl, ctx, -1);
- auto next_token_str = common_token_to_piece(ctx, next_token);
+ // Load state from file
+ std::vector<llama_token> unused_sts(tokens.size());
+ size_t n_token_count_out = 0;
- printf("%s", next_token_str.c_str());
- result0 += next_token_str;
+ 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;
+ }
- common_batch_clear(batch);
- common_batch_add(batch, next_token, n_past, {0}, true);
+ LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
- if (llama_decode(ctx, batch)) {
- fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
- llama_batch_free(batch);
- return 1;
- }
- n_past += 1;
+ // 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;
+}
- printf("\n\n");
- // make new context
- llama_context * ctx2 = llama_init_from_model(model, common_context_params_to_llama(params));
+// 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)};
- llama_sampler * smpl2 = llama_sampler_chain_init(sparams);
+ 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));
- llama_sampler_chain_add(smpl2, llama_sampler_init_dist(params.sampling.seed));
+ auto tokens = common_tokenize(ctx.get(), params.prompt, true);
- printf("\nsecond run: %s", params.prompt.c_str());
+ 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()); // unused session tokens.
+ // Load state from file
+ std::vector<llama_token> unused_sts(tokens.size());
size_t n_token_count_out = 0;
- if (!llama_state_load_file(ctx2, state_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
- fprintf(stderr, "\n%s : failed to load state\n", __func__);
- return 1;
+ 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;
}
- fprintf(stderr, "%s : loaded state with %zu tokens\n", __func__, n_token_count_out);
+ LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
- // restore state (last tokens)
- n_past = n_token_count_out;
- if (!common_replay_last_token(ctx2, tokens.back(), n_past)) {
- return 1;
+ // 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;
+ n_past++;
- // second run
- for (auto i = 0; i < params.n_predict; i++) {
- auto next_token = llama_sampler_sample(smpl2, ctx2, -1);
- auto next_token_str = common_token_to_piece(ctx2, next_token);
-
- printf("%s", next_token_str.c_str());
- result1 += next_token_str;
+ // 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);
- common_batch_clear(batch);
- common_batch_add(batch, next_token, n_past, {0}, true);
+ llama_memory_clear(llama_get_memory(ctx.get()), true);
+ LOG_TRC("%s: kv cache cleared\n", __func__);
- if (llama_decode(ctx2, batch)) {
- fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
- llama_batch_free(batch);
- return 1;
+ 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;
}
- n_past += 1;
+ LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
}
- printf("\n\n");
+ // 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 (result0 != result1) {
- fprintf(stderr, "\n%s : error : the 2 generations are different\n", __func__);
- return 1;
+ if (result != expected_result) {
+ LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
+ return false;
}
- // make new context
- auto params_ctx3 = common_context_params_to_llama(params);
- params_ctx3.n_seq_max = 2;
- llama_context * ctx3 = llama_init_from_model(model, params_ctx3);
+ LOG("\nPASS\n");
+ return true;
+}
+
- llama_sampler * smpl3 = llama_sampler_chain_init(sparams);
+// 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)};
- llama_sampler_chain_add(smpl3, llama_sampler_init_dist(params.sampling.seed));
+ 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));
- printf("\nsingle seq run: %s", params.prompt.c_str());
+ auto tokens = common_tokenize(ctx.get(), params.prompt, true);
- // load state (rng, logits, embedding and kv_cache) from file
- n_token_count_out = 0;
+ LOG("\n=== Test 4: seq copy (device) ===\n");
+ LOG("%s", params.prompt.c_str());
- if (!llama_state_load_file(ctx3, state_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
- fprintf(stderr, "\n%s : failed to load state\n", __func__);
- return 1;
+ // 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;
}
- fprintf(stderr, "%s : loaded state with %zu tokens\n", __func__, n_token_count_out);
+ LOG_TRC("%s: loaded state with %zu tokens\n", __func__, n_token_count_out);
- // restore state (last tokens)
- n_past = n_token_count_out;
- if (!common_replay_last_token(ctx3, tokens.back(), n_past)) {
- return 1;
+ // 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;
+ n_past++;
- // save seq 0 and load into seq 1
+ // Migrate KV cache from seq 0 to seq 1 (on-device path)
{
- // save kv of seq 0
- std::vector<uint8_t> seq_store(llama_state_seq_get_size(ctx3, 0));
- const size_t ncopy = llama_state_seq_get_data(ctx3, seq_store.data(), seq_store.size(), 0);
+ 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()) {
- fprintf(stderr, "\n%s : seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
- return 1;
+ LOG_ERR("\n%s: seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
+ return false;
}
- fprintf(stderr, "%s : seq 0 copied, %zd bytes\n", __func__, ncopy);
+ LOG_TRC("%s: seq 0 copied, %zd bytes\n", __func__, ncopy);
- // erase whole kv
- llama_memory_clear(llama_get_memory(ctx3), true);
- fprintf(stderr, "%s : kv cache cleared\n", __func__);
+ llama_memory_clear(llama_get_memory(ctx.get()), true);
+ LOG_TRC("%s: kv cache cleared\n", __func__);
- // restore kv into seq 1
- const size_t nset = llama_state_seq_set_data(ctx3, seq_store.data(), seq_store.size(), 1);
+ 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()) {
- fprintf(stderr, "\n%s : seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
- return 1;
+ LOG_ERR("\n%s: seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
+ return false;
}
- fprintf(stderr, "%s : seq 1 restored, %zd bytes\n", __func__, nset);
+ LOG_TRC("%s: seq 1 restored, %zd bytes\n", __func__, nset);
}
- // third run with seq 1 instead of 0
- for (auto i = 0; i < params.n_predict; i++) {
- auto next_token = llama_sampler_sample(smpl3, ctx3, -1);
- auto next_token_str = common_token_to_piece(ctx3, next_token);
-
- printf("%s", next_token_str.c_str());
- result2 += next_token_str;
-
- common_batch_clear(batch);
- common_batch_add(batch, next_token, n_past, {1}, true);
+ // 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 (llama_decode(ctx3, batch)) {
- fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
- llama_batch_free(batch);
- return 1;
- }
- n_past += 1;
+ if (result != expected_result) {
+ LOG_ERR("\n%s: error: generation differs from expected\n", __func__);
+ return false;
}
- // test on-device state save/load
- auto params_ctx4 = common_context_params_to_llama(params);
- params_ctx4.n_seq_max = 2;
- llama_context * ctx4 = llama_init_from_model(model, params_ctx4);
+ LOG("\nPASS\n");
+ return true;
+}
- llama_sampler * smpl4 = llama_sampler_chain_init(sparams);
- llama_sampler_chain_add(smpl4, llama_sampler_init_dist(params.sampling.seed));
+int main(int argc, char ** argv) {
+ std::setlocale(LC_NUMERIC, "C");
- printf("\nsingle seq run: %s", params.prompt.c_str());
+ common_params params;
+ params.prompt = "The quick brown fox";
+ params.out_file = "dump_state.bin";
+ params.sampling.seed = 1234;
- // load state (rng, logits, embedding and kv_cache) from file
- n_token_count_out = 0;
+ common_init();
- if (!llama_state_load_file(ctx4, state_file.data(), unused_sts.data(), unused_sts.size(), &n_token_count_out)) {
- fprintf(stderr, "\n%s : failed to load state\n", __func__);
+ if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {
return 1;
}
- fprintf(stderr, "%s : loaded state with %zu tokens\n", __func__, n_token_count_out);
-
- // restore state (last tokens)
- n_past = n_token_count_out;
- if (!common_replay_last_token(ctx4, tokens.back(), n_past)) {
- return 1;
+ if (params.n_parallel == 1) {
+ LOG_TRC("%s: n_parallel == 1, enabling unified kv cache\n", __func__);
+ params.kv_unified = true;
}
- ++n_past;
-
- // save seq 0 and load into seq 1
- {
- // save kv of seq 0
- std::vector<uint8_t> seq_store(llama_state_seq_get_size_ext(ctx4, 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE));
- const size_t ncopy = llama_state_seq_get_data_ext(ctx4, seq_store.data(), seq_store.size(), 0, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
- if (ncopy != seq_store.size()) {
- fprintf(stderr, "\n%s : seq copy data length %zd does not match expected length %zd\n", __func__, ncopy, seq_store.size());
- return 1;
- }
- fprintf(stderr, "%s : seq 0 copied, %zd bytes\n", __func__, ncopy);
-
- // erase whole kv
- llama_memory_clear(llama_get_memory(ctx4), true);
- fprintf(stderr, "%s : kv cache cleared\n", __func__);
- // restore kv into seq 0
- const size_t nset = llama_state_seq_set_data_ext(ctx4, seq_store.data(), seq_store.size(), 1, LLAMA_STATE_SEQ_FLAGS_ON_DEVICE);
- if (nset != seq_store.size()) {
- fprintf(stderr, "\n%s : seq set data length %zd does not match expected length %zd\n", __func__, nset, seq_store.size());
- return 1;
- }
- fprintf(stderr, "%s : seq 1 restored, %zd bytes\n", __func__, nset);
+ if (params.n_predict < 0) {
+ params.n_predict = 16;
}
- // forth run
- for (auto i = 0; i < params.n_predict; i++) {
- auto next_token = llama_sampler_sample(smpl4, ctx4, -1);
- auto next_token_str = common_token_to_piece(ctx4, next_token);
-
- printf("%s", next_token_str.c_str());
- result3 += next_token_str;
+ ggml_backend_load_all();
- common_batch_clear(batch);
- common_batch_add(batch, next_token, n_past, {1}, true);
+ auto llama_init = common_init_from_params(params, true);
+ auto * model = llama_init->model();
- if (llama_decode(ctx4, batch)) {
- fprintf(stderr, "\n%s : failed to evaluate\n", __func__);
- llama_batch_free(batch);
- return 1;
- }
- n_past += 1;
+ if (model == nullptr) {
+ LOG_ERR("%s: failed to init\n", __func__);
+ return 1;
}
- printf("\n");
-
- llama_sampler_free(smpl);
- llama_sampler_free(smpl2);
- llama_sampler_free(smpl3);
- llama_sampler_free(smpl4);
+ GGML_ASSERT(llama_init->context() == nullptr);
- llama_batch_free(batch);
-
- // this one is managed by common_init_result
- //llama_free(ctx);
+ // Test 1: baseline (saves state to disk)
+ auto result_baseline = test_baseline(model, params);
+ if (result_baseline.empty()) {
+ return 1;
+ }
- llama_free(ctx2);
- llama_free(ctx3);
- llama_free(ctx4);
+ // Test 2: state load
+ if (!test_state_load(model, params, result_baseline)) {
+ return 1;
+ }
- if (result0 != result2) {
- fprintf(stderr, "\n%s : error : the seq restore generation is different\n", __func__);
+ // Test 3: seq copy (host)
+ if (!test_seq_cp_host(model, params, result_baseline)) {
return 1;
}
- if (result0 != result3) {
- fprintf(stderr, "\n%s : error : the seq restore generation is different\n", __func__);
+ // Test 4: seq copy (device)
+ if (!test_seq_cp_device(model, params, result_baseline)) {
return 1;
}
- fprintf(stderr, "\n%s : success\n", __func__);
+ LOG("\nAll tests passed.\n");
return 0;
}
overview / pkg / ggml / sources / llama.cpp / commitdiff