bonsai.cpp/examples/rs-rollback/rs-rollback-multi.cpp

// Multi-sequence correctness harness for the recurrent-state snapshot ring.
//
// The ring's snapshot rows are indexed seq-major (see s_write_rows_conv in
// build_conv_state), so a layout error would show up as one sequence's rollback
// corrupting a *different* sequence's state. Single-sequence tests cannot see
// that class of bug at all.
//
//   pass A (reference): run S sequences concurrently, no rollback.
//   pass B (rollback):  identical, but sequence 0 periodically over-generates R
//                       tokens which are then seq_rm()'d away and re-generated.
//
// PASS requires every sequence -- including the untouched ones -- to emit
// exactly the reference token stream.
#include "llama.h"

#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>

static llama_token greedy_ith(llama_context * ctx, const llama_vocab * vocab, int i) {
    const int     n_vocab = llama_vocab_n_tokens(vocab);
    const float * logits  = llama_get_logits_ith(ctx, i);
    llama_token   best    = 0;
    float         best_v  = logits[0];
    for (int t = 1; t < n_vocab; t++) {
        if (logits[t] > best_v) {
            best_v = logits[t];
            best   = t;
        }
    }
    return best;
}

int main(int argc, char ** argv) {
    std::string model_path;
    int         n_predict = 64;
    int         n_seqs    = 3;
    int         stride    = 10;
    int         rewind    = 4;
    int         n_rs_seq  = 8;  // snapshot ring depth; must be >= rewind

    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) {
            n_seqs = 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]);
        }
    }
    if (model_path.empty()) {
        fprintf(stderr, "usage: %s -m model.gguf [-n N] [-S seqs] [-s stride] [-r rewind] [-k n_rs_seq]\n", argv[0]);
        return 1;
    }
    // n_seqs <= 0 would leave the prompt/output vectors empty and run() indexes
    // out[0]; rewind <= 0 would make the rollback branch a no-op
    if (n_seqs <= 0 || rewind <= 0 || stride <= 0 || n_predict <= 0) {
        fprintf(stderr, "invalid arguments: need n_seqs > 0, rewind > 0, stride > 0, n_predict > 0 "
                        "(got %d, %d, %d, %d)\n", n_seqs, rewind, stride, n_predict);
        return 1;
    }

    // deliberately different prompts so the sequences cannot coincidentally agree
    std::vector<std::string> prompts = {
        "Explain how a jet engine works, step by step.",
        "Write a short poem about the sea in winter.",
        "List three causes of the French Revolution.",
        "Describe the process of photosynthesis briefly.",
    };
    if (n_seqs > (int) prompts.size()) {
        n_seqs = (int) prompts.size();
    }

    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);

    std::vector<std::vector<llama_token>> ptoks(n_seqs);
    for (int s = 0; s < n_seqs; s++) {
        const std::string & p = prompts[s];
        int                 n = -llama_tokenize(vocab, p.c_str(), p.size(), NULL, 0, true, true);
        ptoks[s].resize(n);
        llama_tokenize(vocab, p.c_str(), p.size(), ptoks[s].data(), n, true, true);
    }

    int n_rollbacks = 0;

    auto run = [&](bool with_rollback, std::vector<std::vector<llama_token>> & out) -> bool {
        out.assign(n_seqs, {});

        llama_context_params cparams = llama_context_default_params();
        cparams.n_ctx                = (n_predict + 128) * n_seqs;
        cparams.n_batch              = 512;
        cparams.n_seq_max            = n_seqs;
        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_init(512, 0, n_seqs);
        std::vector<int> npos(n_seqs, 0);

        auto submit = [&](const std::vector<std::pair<int, llama_token>> & items) -> bool {
            batch.n_tokens = 0;
            for (auto & it : items) {
                const int s        = it.first;
                const int i        = batch.n_tokens;
                batch.token[i]     = it.second;
                batch.pos[i]       = npos[s];
                batch.n_seq_id[i]  = 1;
                batch.seq_id[i][0] = s;
                batch.logits[i]    = 1;
                batch.n_tokens++;
                npos[s]++;
            }
            return llama_decode(ctx, batch) == 0;
        };

        // prompts: one sequence at a time (each is a multi-token ubatch)
        for (int s = 0; s < n_seqs; s++) {
            batch.n_tokens = 0;
            for (size_t k = 0; k < ptoks[s].size(); k++) {
                batch.token[batch.n_tokens]     = ptoks[s][k];
                batch.pos[batch.n_tokens]       = npos[s];
                batch.n_seq_id[batch.n_tokens]  = 1;
                batch.seq_id[batch.n_tokens][0] = s;
                batch.logits[batch.n_tokens]    = (k + 1 == ptoks[s].size());
                batch.n_tokens++;
                npos[s]++;
            }
            if (llama_decode(ctx, batch)) {
                fprintf(stderr, "decode(prompt %d) failed\n", s);
                return false;
            }
            out[s].push_back(greedy_ith(ctx, vocab, batch.n_tokens - 1));
        }

        int since = 0;
        while ((int) out[0].size() < n_predict) {
            // one token per sequence, all in one batch
            std::vector<std::pair<int, llama_token>> items;
            for (int s = 0; s < n_seqs; s++) {
                items.push_back({ s, out[s].back() });
            }
            if (!submit(items)) {
                fprintf(stderr, "decode(step) failed\n");
                return false;
            }
            for (int s = 0; s < n_seqs; s++) {
                out[s].push_back(greedy_ith(ctx, vocab, s));
            }
            since++;

            if (with_rollback && since >= stride && (int) out[0].size() + rewind + 2 < n_predict) {
                n_rollbacks++;
                // seq 0 alone over-generates, then discards
                for (int r = 0; r < rewind; r++) {
                    std::vector<std::pair<int, llama_token>> one = {
                        { 0, out[0].back() }
                    };
                    if (!submit(one)) {
                        fprintf(stderr, "decode(draft) failed\n");
                        return false;
                    }
                    out[0].push_back(greedy_ith(ctx, vocab, 0));
                }
                for (int r = 0; r < rewind; r++) {
                    out[0].pop_back();
                }

                // invariant: out[s].back() is sampled-but-not-yet-decoded. The
                // drafts occupied positions [npos0 - rewind, npos0), and that
                // range also covers the decode of the still-pending token, so
                // discarding it restores the exact pre-draft state -- no replay.
                const llama_pos keep = npos[0] - rewind;
                if (!llama_memory_seq_rm(mem, 0, keep, -1)) {
                    fprintf(stderr, "ROLLBACK REFUSED: keep=%d npos0=%d rewind=%d\n", keep, npos[0], rewind);
                    return false;
                }
                npos[0] = keep;
                since   = 0;
            }
        }

        llama_batch_free(batch);
        llama_free(ctx);
        return true;
    };

    std::vector<std::vector<llama_token>> ref, rb;
    if (!run(false, ref)) {
        return 1;
    }
    if (!run(true, rb)) {
        return 1;
    }

    printf("\n=== rs-rollback-multi (S=%d, rewind=%d) ===\n", n_seqs, rewind);
    int total_mismatch = 0;
    for (int s = 0; s < n_seqs; s++) {
        size_t n  = ref[s].size() < rb[s].size() ? ref[s].size() : rb[s].size();
        int    mm = 0, first = -1;
        for (size_t i = 0; i < n; i++) {
            if (ref[s][i] != rb[s][i]) {
                if (first < 0) {
                    first = (int) i;
                }
                mm++;
            }
        }
        total_mismatch += mm;
        printf("  seq %d %-14s compared=%zu mismatches=%d%s\n", s, s == 0 ? "(rolled back)" : "(untouched)", n, mm,
               mm ? "" : "  OK");
        if (mm) {
            printf("      first mismatch at %d (ref=%d rb=%d)\n", first, ref[s][first], rb[s][first]);
        }
    }
    printf("rollbacks done: %d\n", n_rollbacks);
    // no rollback attempted means the two runs were identical by construction
    const bool ok = (n_rollbacks > 0) && (total_mismatch == 0);
    if (n_rollbacks == 0) {
        printf("RESULT: FAIL (no rollback was attempted: need n_predict > stride + rewind + 2, "
               "got n_predict=%d stride=%d rewind=%d)\n", n_predict, stride, rewind);
    } else {
        printf("RESULT: %s\n", total_mismatch ? "FAIL" : "PASS");
    }

    llama_model_free(model);
    llama_backend_free();
    return ok ? 0 : 1;
}