86 lines
3.3 KiB
C++
86 lines
3.3 KiB
C++
#include "llama.h"
|
|
#include "ggml-backend.h"
|
|
|
|
#include <cstdio>
|
|
#include <cstdlib>
|
|
#include <cstring>
|
|
#include <fstream>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
static void require(bool ok, const char * what) {
|
|
if (!ok) {
|
|
std::fprintf(stderr, "FAIL: %s\n", what);
|
|
std::exit(1);
|
|
}
|
|
}
|
|
|
|
static std::vector<uint8_t> save(llama_context * ctx) {
|
|
std::vector<uint8_t> data(llama_state_seq_get_size(ctx, 0));
|
|
require(llama_state_seq_get_data(ctx, data.data(), data.size(), 0) == data.size(), "save");
|
|
return data;
|
|
}
|
|
|
|
int main(int argc, char ** argv) {
|
|
if (argc != 3) return 2;
|
|
ggml_backend_load_all();
|
|
llama_backend_init();
|
|
auto mp = llama_model_default_params();
|
|
mp.n_gpu_layers = 999;
|
|
auto model = llama_model_load_from_file(argv[1], mp);
|
|
require(model != nullptr, "model");
|
|
auto cp = llama_context_default_params();
|
|
cp.n_ctx = 1024;
|
|
cp.n_batch = 512;
|
|
cp.n_ubatch = 512;
|
|
cp.n_seq_max = 1;
|
|
cp.n_rs_seq = 3;
|
|
cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
|
|
auto ctx = llama_init_from_model(model, cp);
|
|
require(ctx != nullptr, "context");
|
|
auto mem = llama_get_memory(ctx);
|
|
auto batch = llama_batch_init(512, 0, 1);
|
|
const int vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
|
|
auto decode = [&](int start, int count) {
|
|
batch.n_tokens = count;
|
|
for (int i = 0; i < count; ++i) {
|
|
batch.token[i] = 100 + (start + i) % 1000;
|
|
batch.pos[i] = start + i;
|
|
batch.n_seq_id[i] = 1;
|
|
batch.seq_id[i][0] = 0;
|
|
batch.logits[i] = i == count - 1;
|
|
}
|
|
require(llama_decode(ctx, batch) == 0, "decode");
|
|
};
|
|
auto record = [&](const std::string & name, const std::vector<uint8_t> & data) {
|
|
std::ofstream file(std::string(argv[2]) + "/" + name, std::ios::binary);
|
|
file.write((const char *) data.data(), data.size());
|
|
require(file.good(), "write");
|
|
};
|
|
decode(0, 32);
|
|
auto base = save(ctx);
|
|
record("base.bin", base);
|
|
for (int keep = 1; keep <= 4; ++keep) {
|
|
require(llama_state_seq_set_data(ctx, base.data(), base.size(), 0) == base.size(), "restore base");
|
|
decode(32, 4);
|
|
if (keep < 4) require(llama_memory_seq_rm(mem, 0, 32 + keep, -1), "rollback");
|
|
auto cut = save(ctx);
|
|
record("cut-" + std::to_string(keep) + ".bin", cut);
|
|
if (keep < 4) require(!llama_memory_seq_rm(mem, 0, 32 + keep - 1, -1), "reject second pending rollback");
|
|
decode(32 + keep, 1);
|
|
std::vector<float> logits(llama_get_logits_ith(ctx, -1), llama_get_logits_ith(ctx, -1) + vocab);
|
|
require(llama_state_seq_set_data(ctx, cut.data(), cut.size(), 0) == cut.size(), "restore cut");
|
|
decode(32 + keep, 1);
|
|
require(memcmp(logits.data(), llama_get_logits_ith(ctx, -1), vocab * sizeof(float)) == 0, "save/load logits parity");
|
|
record("after-" + std::to_string(keep) + ".bin", save(ctx));
|
|
}
|
|
require(llama_memory_seq_rm(mem, 0, 0, -1), "clear sequence");
|
|
decode(0, 32);
|
|
require(save(ctx) == base, "clear and restart");
|
|
llama_batch_free(batch);
|
|
llama_free(ctx);
|
|
llama_model_free(model);
|
|
llama_backend_free();
|
|
std::puts("PASS: accepted-prefix rollback, pending rollback rejection, state save/load, clear and restart");
|
|
}
|