// Correctness harness for the recurrent-state snapshot ring (n_rs_seq > 0).
//
// A forward-only generation never reads the snapshot slots back, so it cannot
// catch a bad snapshot write. This does:
//
// pass A (reference): greedily generate N tokens straight through.
// pass B (rollback): generate the same N tokens, but every S tokens
// over-generate by R, then llama_memory_seq_rm() those R
// positions away and re-generate them.
//
// Pass B only lands on the same tokens if the state restored from the snapshot
// ring is exactly the state that produced them the first time. Any error in the
// conv-window or SSM snapshot write shows up as a token mismatch.
#include "llama.h"
#include <cstdio>
#include <cstring>
#include <string>
#include <vector>
static llama_token greedy(llama_context * ctx, const llama_vocab * vocab) {
const int n_vocab = llama_vocab_n_tokens(vocab);
const float * logits = llama_get_logits_ith(ctx, -1);
llama_token best = 0;
float best_v = logits[0];
for (int i = 1; i < n_vocab; i++) {
if (logits[i] > best_v) {
best_v = logits[i];
best = i;
}
}
return best;
}
int main(int argc, char ** argv) {
std::string model_path;
std::string prompt = "Explain how a jet engine works, step by step.";
int n_predict = 128;
int stride = 16; // roll back every `stride` accepted tokens
int rewind = 4; // how many tokens to discard and re-generate
int n_rs_seq = 8; // snapshot ring depth; must be >= rewind + 1
for (int i = 1; i < argc; i++) {
if (!strcmp(argv[i], "-m") && i + 1 < argc) {
model_path = argv[++i];
} else if (!strcmp(argv[i], "-n") && i + 1 < argc) {
n_predict = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-s") && i + 1 < argc) {
stride = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-r") && i + 1 < argc) {
rewind = atoi(argv[++i]);
} else if (!strcmp(argv[i], "-k") && i + 1 < argc) {
n_rs_seq = atoi(argv[++i]);
} else if (argv[i][0] != '-') {
prompt = argv[i];
}
}
if (model_path.empty()) {
fprintf(stderr, "usage: %s -m model.gguf [-n N] [-s stride] [-r rewind] [-k n_rs_seq]\n", argv[0]);
return 1;
}
llama_backend_init();
llama_model_params mparams = llama_model_default_params();
mparams.n_gpu_layers = 99;
llama_model * model = llama_model_load_from_file(model_path.c_str(), mparams);
if (!model) {
fprintf(stderr, "failed to load model\n");
return 1;
}
const llama_vocab * vocab = llama_model_get_vocab(model);
const int n_prompt = -llama_tokenize(vocab, prompt.c_str(), prompt.size(), NULL, 0, true, true);
std::vector<llama_token> ptoks(n_prompt);
llama_tokenize(vocab, prompt.c_str(), prompt.size(), ptoks.data(), ptoks.size(), true, true);
int n_rollbacks = 0;
auto run = [&](bool with_rollback, std::vector<llama_token> & out) -> bool {
llama_context_params cparams = llama_context_default_params();
cparams.n_ctx = n_prompt + n_predict + 64;
cparams.n_batch = n_prompt + 64;
cparams.n_rs_seq = n_rs_seq;
llama_context * ctx = llama_init_from_model(model, cparams);
if (!ctx) {
fprintf(stderr, "failed to create context\n");
return false;
}
llama_memory_t mem = llama_get_memory(ctx);
llama_batch batch = llama_batch_get_one(ptoks.data(), ptoks.size());
if (llama_decode(ctx, batch)) {
fprintf(stderr, "decode(prompt) failed\n");
return false;
}
int n_pos = n_prompt;
int since = 0;
while ((int) out.size() < n_predict) {
llama_token tok = greedy(ctx, vocab);
out.push_back(tok);
batch = llama_batch_get_one(&tok, 1);
if (llama_decode(ctx, batch)) {
fprintf(stderr, "decode failed\n");
return false;
}
n_pos++;
since++;
if (with_rollback && since >= stride && (int) out.size() + rewind < n_predict) {
n_rollbacks++;
// over-generate `rewind` tokens, then throw them away
for (int r = 0; r < rewind; r++) {
llama_token t = greedy(ctx, vocab);
batch = llama_batch_get_one(&t, 1);
if (llama_decode(ctx, batch)) {
fprintf(stderr, "decode(draft) failed\n");
return false;
}
n_pos++;
}
// Drop the drafts AND the last accepted token in one go, then
// re-decode the accepted token so the logits are valid again.
// This forces the state to be restored from a snapshot taken
// (rewind + 1) steps back -- needs n_rs_seq >= rewind + 1.
const llama_pos keep = n_pos - rewind - 1;
if (!llama_memory_seq_rm(mem, 0, keep, -1)) {
fprintf(stderr, "ROLLBACK REFUSED: keep=%d n_pos=%d rewind=%d\n", keep, n_pos, rewind);
return false;
}
n_pos = keep;
batch = llama_batch_get_one(&out.back(), 1);
if (llama_decode(ctx, batch)) {
fprintf(stderr, "decode(replay) failed\n");
return false;
}
n_pos++;
since = 0;
}
}
llama_free(ctx);
return true;
};
bool ok = true;
std::vector<llama_token> ref, rb;
if (!run(false, ref)) {
return 1;
}
if (!run(true, rb)) {
return 1;
}
size_t n = ref.size() < rb.size() ? ref.size() : rb.size();
size_t mismatch = 0;
int first = -1;
for (size_t i = 0; i < n; i++) {
if (ref[i] != rb[i]) {
if (first < 0) {
first = (int) i;
}
mismatch++;
}
}
printf("\n=== rs-rollback ===\n");
printf("tokens compared : %zu\n", n);
printf("mismatches : %zu\n", mismatch);
printf("rollbacks done : %d\n", n_rollbacks);
// a run that never entered the rollback branch compared two identical
// straight-through generations and proved nothing, so it is not a pass
if (n_rollbacks == 0) {
printf("RESULT: FAIL (no rollback was attempted: need n_predict > stride + rewind, "
"got n_predict=%d stride=%d rewind=%d)\n", n_predict, stride, rewind);
ok = false;
} else if (mismatch) {
printf("first mismatch : index %d (ref=%d rb=%d)\n", first, ref[first], rb[first]);
printf("RESULT: FAIL\n");
ok = false;
} else {
printf("RESULT: PASS (rollback reproduces the straight-through sequence)\n");
}
llama_model_free(model);
llama_backend_free();
return ok ? 0 : 1;
}