Something went wrong. Try again.
llama.cpp fork for Ternary Bonsai 2 27B: MTP speculation, suffix cache reuse (SCR), q8 KV, two-slot fairness, fused DeltaNet kernels
Something went wrong. Try again.
15 kB · 440 lines
C++
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441#include "arg.h"#include "common.h"#include "log.h"#include "llama-cpp.h"
#include <clocale>#include <random>#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 llama_tokens generate_tokens(llama_context * ctx, llama_sampler * smpl, int & n_past, int32_t n_predict, llama_seq_id seq_id) { llama_tokens 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);
LOG("%d ", next_token); result.push_back(next_token);
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// - decode all but the last token// - save state to disk// - decode the last token// - generate n_predict tokensstatic llama_tokens test_baseline(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { 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 n_past = 0; if (!common_prompt_batch_decode(ctx.get(), tokens, (int)tokens.size(), 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");
auto result = generate_tokens(ctx.get(), smpl.get(), n_past, params.n_predict, 0); if (result.empty()) { return {}; }
LOG("\n");
return result;}
// Test 2: sequence removal isolation// - decode the same prefix into two sequences// - remove sequence 0// - verify that sequence 1 remains unchangedstatic bool test_seq_rm_isolated( struct llama_model * model, const struct common_params & params, const llama_tokens & tokens) { auto params_ctx = common_context_params_to_llama(params); params_ctx.n_ctx = 256; params_ctx.n_seq_max = 2; params_ctx.kv_unified = true;
auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; if (!ctx) { LOG_ERR("%s: failed to create context\n", __func__); return false; }
LOG("\n=== Test 2: sequence removal isolation ===\n");
const size_t n_tokens = tokens.size() < 128 ? tokens.size() : 128; for (llama_seq_id seq_id = 0; seq_id < 2; ++seq_id) { llama_batch_ptr batch(n_tokens, 0, 1); for (size_t i = 0; i < n_tokens; ++i) { common_batch_add(batch.get(), tokens[i], i, { seq_id }, false); }
if (llama_decode(ctx.get(), batch.get())) { LOG_ERR("%s: failed to decode prompt for sequence %d\n", __func__, seq_id); return false; } }
const auto get_seq_state = [&](llama_seq_id seq_id, std::vector<uint8_t> & state) { const size_t state_size = llama_state_seq_get_size(ctx.get(), seq_id); if (state_size == 0) { LOG_ERR("%s: sequence state is empty\n", __func__); return false; }
state.resize(state_size); const size_t ncopy = llama_state_seq_get_data(ctx.get(), state.data(), state.size(), seq_id); if (ncopy != state.size()) { LOG_ERR("%s: sequence state length %zu does not match expected length %zu\n", __func__, ncopy, state.size()); return false; }
return true; };
std::vector<uint8_t> state_before; if (!get_seq_state(1, state_before)) { return false; }
if (!llama_memory_seq_rm(llama_get_memory(ctx.get()), 0, -1, -1)) { LOG_ERR("%s: failed to remove sequence 0\n", __func__); return false; }
std::vector<uint8_t> state_after; if (!get_seq_state(1, state_after)) { return false; }
if (state_before != state_after) { LOG_ERR("%s: removing sequence 0 changed sequence 1\n", __func__); return false; }
LOG("PASS\n"); return true;}
// Test 3: state load// - create a new context// - load state from file// - replay the last prompt token// - generate n_predict tokens and compare against expected resultstatic bool test_state_load(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & 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));
LOG("\n=== Test 3: state load ===\n");
// Load state from file llama_tokens 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 - 1; 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 4: 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 resultstatic bool test_seq_cp_host(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & 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));
LOG("\n=== Test 4: seq copy (host) ===\n");
// Load state from file llama_tokens 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 - 1; 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 5: 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 resultstatic bool test_seq_cp_device(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, const llama_tokens & 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));
LOG("\n=== Test 5: seq copy (device) ===\n");
// Load state from file llama_tokens 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 - 1; 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 = ""; params.n_batch = 100; 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);
// Tokenize prompt or generate random tokens llama_tokens tokens; if (params.prompt.empty()) { const int n_prompt = params.n_batch;
// this path is useful for model files that do not have a tokenizer LOG_INF("%s: no prompt provided, generating %d (n_batch) random tokens\n", __func__, n_prompt);
const auto * vocab = llama_model_get_vocab(model); const auto n_vocab = llama_vocab_n_tokens(vocab);
std::mt19937 rng(params.sampling.seed); std::uniform_int_distribution<llama_token> dist(0, n_vocab - 1); for (int i = 0; i < n_prompt; i++) { tokens.push_back(dist(rng)); } } else { LOG_INF("%s: tokenizing prompt '%s'\n", __func__, params.prompt.c_str());
auto ctx = llama_context_ptr{llama_init_from_model(model, common_context_params_to_llama(params))}; tokens = common_tokenize(ctx.get(), params.prompt, true); }
LOG_INF("%s: the input prompt is %d tokens\n", __func__, (int)tokens.size());
// Test 1: baseline (saves state to disk) auto result_baseline = test_baseline(model, params, tokens); if (result_baseline.empty()) { return 1; }
// Test 2: sequence removal isolation if (!test_seq_rm_isolated(model, params, tokens)) { return 1; }
// Test 3: state load if (!test_state_load(model, params, tokens, result_baseline)) { return 1; }
// Test 4: seq copy (host) if (!test_seq_cp_host(model, params, tokens, result_baseline)) { return 1; }
// Test 5: seq copy (device) if (!test_seq_cp_device(model, params, tokens, result_baseline)) { return 1; }
LOG("\nAll tests passed.\n");
return 0;}