Felipe97/llama-cpp-compiled

1#include "arg.h"2#include "common.h"3#include "ggml-backend.h"4#include "llama.h"5 6#include "../src/llama-io.h"7#include "../src/llama-memory.h"8 9#include <algorithm>10#include <clocale>11#include <cmath>12#include <cstdio>13#include <limits>14#include <set>15#include <vector>16 17static bool decode_tokens(llama_context * ctx, const std::vector<llama_token> & tokens, uint32_t count) {18    llama_batch batch = llama_batch_init(count, 0, 1);19    for (uint32_t pos = 0; pos < count; ++pos) {20        common_batch_add(batch, tokens[pos], pos, { 0 }, pos + 1 == count);21    }22    const bool ok = llama_decode(ctx, batch) == 0;23    llama_batch_free(batch);24    return ok;25}26 27static bool decode_one(llama_context * ctx, llama_token tok, llama_pos pos) {28    llama_batch batch = llama_batch_init(1, 0, 1);29    common_batch_add(batch, tok, pos, { 0 }, true);30    const bool ok = llama_decode(ctx, batch) == 0;31    llama_batch_free(batch);32    return ok;33}34 35struct cache_buffer_collector : llama_io_write_i {36    std::set<ggml_backend_buffer_t> buffers;37    size_t size = 0;38 39    void write(const void *, size_t n) override {40        size += n;41    }42 43    void write_tensor(ggml_tensor * tensor, size_t, size_t n) override {44        buffers.insert(tensor->buffer);45        size += n;46    }47 48    size_t n_bytes() override {49        return size;50    }51};52 53static llama_context * init_ctx(llama_model * model, llama_context_params cparams, uint8_t fill) {54    llama_context * ctx = llama_init_from_model(model, cparams);55    if (ctx == nullptr || fill == 0) {56        return ctx;57    }58 59    // Use a full ubatch so buffer discovery preserves prefill allocation sizes.60    const uint32_t n_tokens = llama_n_ubatch(ctx);61    if (!decode_tokens(ctx, std::vector<llama_token>(n_tokens, 0), n_tokens)) {62        llama_free(ctx);63        return nullptr;64    }65    llama_synchronize(ctx);66    cache_buffer_collector collector;67    llama_get_memory(ctx)->state_write(collector);68    llama_memory_clear(llama_get_memory(ctx), true);69    if (collector.buffers.empty()) {70        fprintf(stderr, "%s : no cache buffers found\n", __func__);71        llama_free(ctx);72        return nullptr;73    }74    for (auto * buffer : collector.buffers) {75        ggml_backend_buffer_clear(buffer, fill);76    }77    return ctx;78}79 80static llama_context * make_ctx(const common_params & params, llama_model * model, uint8_t fill) {81    auto cparams = common_context_params_to_llama(params);82    cparams.n_seq_max = 1;83    cparams.n_rs_seq  = 8;84    cparams.n_batch   = std::max(cparams.n_batch,  (uint32_t) (cparams.n_rs_seq + 1));85    cparams.n_ubatch  = std::max(cparams.n_ubatch, (uint32_t) (cparams.n_rs_seq + 1));86    return init_ctx(model, cparams, fill);87}88 89static float logit_diff(float a, float b) {90    return std::isfinite(a) && std::isfinite(b) ? std::fabs(a - b) : std::numeric_limits<float>::infinity();91}92 93// Roll back multiple sequences, then replay them in a single batch whose94// per-seq token count exceeds n_ubatch: each seq's replay spans several95// ubatches while its rollback restore is still pending. Compared against a96// reference context that never advanced past the rollback point and decodes97// the identical replay batch.98static bool test_multi_seq_split_replay(const common_params & params, llama_model * model, const int n_vocab, uint8_t fill) {99    constexpr uint32_t  n_seqs     = 2;100    constexpr uint32_t  n_ubatch   = 16;101    constexpr uint32_t  n_prompt   = 19;102    constexpr uint32_t  n_rollback = 3;103    constexpr uint32_t  n_replay   = 40; // > n_ubatch so each seq spans multiple ubatches104    constexpr llama_pos p0         = n_prompt - n_rollback;105 106    const auto make_ctx_multi = [&]() {107        auto cparams = common_context_params_to_llama(params);108        cparams.n_seq_max  = n_seqs;109        cparams.n_rs_seq   = 8;110        cparams.n_ctx      = 256;111        cparams.n_batch    = 256;112        cparams.n_ubatch   = n_ubatch;113        cparams.kv_unified = false;114        return init_ctx(model, cparams, fill);115    };116 117    llama_context * ctx_roll = make_ctx_multi();118    llama_context * ctx_ref  = make_ctx_multi();119    if (ctx_roll == nullptr || ctx_ref == nullptr) {120        fprintf(stderr, "%s : failed to init multi-seq contexts\n", __func__);121        return false;122    }123 124    const auto cleanup = [&]() {125        llama_free(ctx_roll);126        llama_free(ctx_ref);127    };128 129    if (llama_n_rs_seq(ctx_roll) < n_rollback) {130        fprintf(stderr, "%s : skipping because n_rs_seq is too small\n", __func__);131        cleanup();132        return true;133    }134 135    const auto tok = [&](uint32_t seq, llama_pos pos) {136        return (llama_token) ((7*(uint32_t) pos + 31*seq + 1) % (uint32_t) n_vocab);137    };138 139    bool ok = true;140 141    // both contexts decode the identical [0, p0) prefill; only ctx_roll decodes142    // the tail, which is then rolled back so its restore is pending at replay143    for (uint32_t s = 0; s < n_seqs && ok; ++s) {144        llama_batch batch = llama_batch_init(n_prompt, 0, 1);145        for (llama_pos pos = 0; pos < (llama_pos) p0; ++pos) {146            common_batch_add(batch, tok(s, pos), pos, { (llama_seq_id) s }, false);147        }148        ok = ok && llama_decode(ctx_roll, batch) == 0;149        ok = ok && llama_decode(ctx_ref,  batch) == 0;150 151        common_batch_clear(batch);152        for (llama_pos pos = p0; pos < (llama_pos) n_prompt; ++pos) {153            common_batch_add(batch, tok(s, pos), pos, { (llama_seq_id) s }, false);154        }155        ok = ok && llama_decode(ctx_roll, batch) == 0;156        llama_batch_free(batch);157 158        ok = ok && llama_memory_seq_rm(llama_get_memory(ctx_roll), (llama_seq_id) s, p0, -1);159 160        // a second partial removal while one is pending must be refused161        ok = ok && !llama_memory_seq_rm(llama_get_memory(ctx_roll), (llama_seq_id) s, p0 - 1, -1);162    }163    if (!ok) {164        fprintf(stderr, "%s : multi-seq prefill/rollback failed\n", __func__);165        cleanup();166        return false;167    }168 169    llama_batch batch = llama_batch_init(n_seqs*n_replay, 0, 1);170    for (uint32_t s = 0; s < n_seqs; ++s) {171        for (uint32_t i = 0; i < n_replay; ++i) {172            const llama_pos pos = p0 + (llama_pos) i;173            common_batch_add(batch, tok(s, pos), pos, { (llama_seq_id) s }, true);174        }175    }176    ok = llama_decode(ctx_roll, batch) == 0;177    ok = ok && llama_decode(ctx_ref, batch) == 0;178    llama_batch_free(batch);179    if (!ok) {180        fprintf(stderr, "%s : multi-seq replay decode failed\n", __func__);181        cleanup();182        return false;183    }184 185    // identical ubatch shapes from bit-exact states: a correct implementation186    // matches bitwise, so eps only allows backend scheduling noise187    constexpr float eps = 1e-7f;188 189    float    diff_max  = 0.0f;190    uint32_t seq_first = 0;191    int32_t  pos_first = -1;192    for (uint32_t i = 0; i < n_seqs*n_replay; ++i) {193        const float * l_roll = llama_get_logits_ith(ctx_roll, i);194        const float * l_ref  = llama_get_logits_ith(ctx_ref,  i);195        if (l_roll == nullptr || l_ref == nullptr) {196            fprintf(stderr, "%s : missing multi-seq logits at index %u\n", __func__, i);197            cleanup();198            return false;199        }200        for (int t = 0; t < n_vocab; ++t) {201            const float diff = logit_diff(l_roll[t], l_ref[t]);202            if (diff > eps && pos_first < 0) {203                seq_first = i/n_replay;204                pos_first = p0 + (int32_t) (i%n_replay);205            }206            diff_max = std::max(diff_max, diff);207        }208    }209 210    if (diff_max > eps) {211        fprintf(stderr, "%s : multi-seq split replay logits mismatch (max diff %g, first at seq %u pos %d)\n",212                __func__, (double) diff_max, seq_first, pos_first);213        cleanup();214        return false;215    }216 217    fprintf(stderr, "%s : multi-seq split replay matched (max diff %g)\n", __func__, (double) diff_max);218 219    // seq-1-only decodes must be independent of seq 0's content: diverge seq 0220    // in ctx_ref only, then compare identical seq-1-only continuations bitwise221    constexpr uint32_t n_tail = 4;222 223    {224        llama_batch batch_tail = llama_batch_init(n_tail, 0, 1);225        for (uint32_t i = 0; i < n_tail; ++i) {226            const llama_pos pos = p0 + (llama_pos) (n_replay + i);227            common_batch_add(batch_tail, tok(0, pos + 7), pos, { 0 }, false);228        }229        ok = llama_decode(ctx_ref, batch_tail) == 0;230        llama_batch_free(batch_tail);231    }232 233    float diff_tail = 0.0f;234    for (uint32_t i = 0; i < n_tail && ok; ++i) {235        const llama_pos pos = p0 + (llama_pos) (n_replay + i);236        llama_batch batch_one = llama_batch_init(1, 0, 1);237        common_batch_add(batch_one, tok(1, pos), pos, { 1 }, true);238        ok = llama_decode(ctx_roll, batch_one) == 0;239        ok = ok && llama_decode(ctx_ref, batch_one) == 0;240        llama_batch_free(batch_one);241        if (!ok) {242            break;243        }244 245        const float * l_roll = llama_get_logits_ith(ctx_roll, 0);246        const float * l_ref  = llama_get_logits_ith(ctx_ref,  0);247        ok = l_roll != nullptr && l_ref != nullptr;248        for (int t = 0; ok && t < n_vocab; ++t) {249            diff_tail = std::max(diff_tail, logit_diff(l_roll[t], l_ref[t]));250        }251    }252 253    if (!ok || diff_tail > eps) {254        fprintf(stderr, "%s : seq-1-only decode leaked seq 0 state (ok=%d, max diff %g)\n",255                __func__, ok ? 1 : 0, (double) diff_tail);256        cleanup();257        return false;258    }259 260    fprintf(stderr, "%s : seq-1-only decode independent of seq 0 (max diff %g)\n", __func__, (double) diff_tail);261    cleanup();262    return true;263}264 265static int test_rollback(const common_params & params, llama_model * model, uint8_t fill) {266    const llama_vocab * vocab   = llama_model_get_vocab(model);267    const int           n_vocab = llama_vocab_n_tokens(vocab);268 269    llama_context * ctx_src = make_ctx(params, model, fill);270    llama_context * ctx_dst = make_ctx(params, model, fill);271    if (ctx_src == nullptr || ctx_dst == nullptr) {272        fprintf(stderr, "%s : failed to init contexts\n", __func__);273        return 1;274    }275 276    if (llama_n_rs_seq(ctx_src) == 0) {277        fprintf(stderr, "%s : skipping because n_rs_seq is disabled\n", __func__);278        llama_free(ctx_src);279        llama_free(ctx_dst);280        return 0;281    }282 283    std::vector<llama_token> tokens;284    if (llama_vocab_type(vocab) == LLAMA_VOCAB_TYPE_NONE) {285        tokens = { 1, 2, 3, 4, 5, 6, 7, 8, 9 };286    } else {287        tokens = common_tokenize(ctx_src, "The quick brown fox jumps over the lazy dog", true);288    }289    const uint32_t n_rs_seq = llama_n_rs_seq(ctx_src);290    constexpr uint32_t n_rollback = 3;291    if (n_rs_seq < n_rollback) {292        fprintf(stderr, "%s : skipping because n_rs_seq is too small\n", __func__);293        llama_free(ctx_src);294        llama_free(ctx_dst);295        return 0;296    }297    if (tokens.empty()) {298        fprintf(stderr, "%s : not enough prompt tokens\n", __func__);299        return 1;300    }301    tokens.resize(n_rs_seq + 1, tokens.back());302 303    const uint32_t  n_tokens     = tokens.size();304    const llama_pos rollback_pos = (llama_pos) n_tokens - n_rollback;305 306    // Decode the full prompt on the source, then roll back three positions.307    // Replaying them crosses DSV4's ratio-4 compressor boundary.308    // Rollback leaves the recurrent memory in a snapshot state (rs_idx != 0).309    if (!decode_tokens(ctx_src, tokens, n_tokens)) {310        fprintf(stderr, "%s : failed to decode prompt\n", __func__);311        return 1;312    }313    if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1)) {314        fprintf(stderr, "%s : rollback failed\n", __func__);315        return 1;316    }317 318    // Save the rolled-back state and restore it into a fresh context.319    common_prompt_checkpoint ckpt;320    ckpt.update_tgt(ctx_src, 0, 0);321    ckpt.load_tgt(ctx_dst, 0, 0);322 323    constexpr float eps = 1e-5f;324    std::vector<std::vector<float>> logits_src_replay(n_rollback);325    const auto replay_and_compare = [&](const char * mode) {326        for (uint32_t i = 0; i < n_rollback; ++i) {327            const llama_pos pos = rollback_pos + i;328            if (!decode_one(ctx_src, tokens[pos], pos) ||329                !decode_one(ctx_dst, tokens[pos], pos)) {330                fprintf(stderr, "%s : %s replay failed at position %d\n", __func__, mode, pos);331                return false;332            }333 334            const float * logits_src = llama_get_logits_ith(ctx_src, 0);335            const float * logits_dst = llama_get_logits_ith(ctx_dst, 0);336            if (logits_src == nullptr || logits_dst == nullptr) {337                fprintf(stderr, "%s : missing %s logits at position %d\n", __func__, mode, pos);338                return false;339            }340 341            logits_src_replay[i].assign(logits_src, logits_src + n_vocab);342            for (int token = 0; token < n_vocab; ++token) {343                if (logit_diff(logits_src[token], logits_dst[token]) > eps) {344                    fprintf(stderr, "%s : %s logits mismatch at position %d, token %d (%g != %g)\n",345                            __func__, mode, pos, token, (double) logits_src[token], (double) logits_dst[token]);346                    return false;347                }348            }349        }350        return true;351    };352    if (!replay_and_compare("full")) {353        return 1;354    }355 356    if (!llama_memory_seq_rm(llama_get_memory(ctx_src), 0, rollback_pos, -1) ||357        !llama_memory_seq_rm(llama_get_memory(ctx_dst), 0, rollback_pos, -1)) {358        fprintf(stderr, "%s : partial rollback failed\n", __func__);359        return 1;360    }361 362    constexpr llama_state_seq_flags partial_flags = LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY;363    common_prompt_checkpoint ckpt_partial;364    ckpt_partial.update_tgt(ctx_src, 0, partial_flags);365    ckpt_partial.load_tgt(ctx_dst, 0, partial_flags);366 367    if (!replay_and_compare("partial")) {368        return 1;369    }370 371    // Repeat the load into a context that already has its own rollback state:372    // groups 1..n_rs_seq hold a different prompt's history, and rs_idx[0] is373    // non-zero at load time. The restore must wipe that state and still match.374    llama_context * ctx_dirty = make_ctx(params, model, fill);375    if (ctx_dirty == nullptr) {376        fprintf(stderr, "%s : failed to init dirty ctx\n", __func__);377        return 1;378    }379 380    std::vector<llama_token> noise = tokens;381    for (auto & t : noise) {382        t = (t + 1) % n_vocab;383        if (t < 0) {384            t = 0;385        }386    }387    if (!decode_tokens(ctx_dirty, noise, n_tokens)) {388        fprintf(stderr, "%s : dirty prompt decode failed\n", __func__);389        return 1;390    }391    if (!llama_memory_seq_rm(llama_get_memory(ctx_dirty), 0, rollback_pos, -1)) {392        fprintf(stderr, "%s : dirty rollback failed\n", __func__);393        return 1;394    }395 396    ckpt.load_tgt(ctx_dirty, 0, 0);397 398    for (uint32_t i = 0; i < n_rollback; ++i) {399        const llama_pos pos = rollback_pos + i;400        if (!decode_one(ctx_dirty, tokens[pos], pos)) {401            fprintf(stderr, "%s : dirty replay failed at position %d\n", __func__, pos);402            return 1;403        }404 405        const float * logits_dirty = llama_get_logits_ith(ctx_dirty, 0);406        if (logits_dirty == nullptr) {407            fprintf(stderr, "%s : missing dirty logits at position %d\n", __func__, pos);408            return 1;409        }410 411        for (int token = 0; token < n_vocab; ++token) {412            if (logit_diff(logits_src_replay[i][token], logits_dirty[token]) > eps) {413                fprintf(stderr, "%s : dirty-ctx logits mismatch at position %d, token %d (%g != %g)\n",414                        __func__, pos, token, (double) logits_src_replay[i][token], (double) logits_dirty[token]);415                return 1;416            }417        }418    }419 420    fprintf(stderr, "%s : recurrent rollback checkpoint restored successfully\n", __func__);421    llama_free(ctx_src);422    llama_free(ctx_dst);423    llama_free(ctx_dirty);424 425    if (!test_multi_seq_split_replay(params, model, n_vocab, fill)) {426        return 1;427    }428 429    return 0;430}431 432int main(int argc, char ** argv) {433    std::setlocale(LC_NUMERIC, "C");434 435    common_params params;436    params.sampling.seed = 1234;437    params.n_predict = 1;438 439    common_init();440 441    if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) {442        return 1;443    }444 445    ggml_backend_load_all();446 447    common_init_result_ptr llama_init = common_init_from_params(params);448    llama_model * model = llama_init->model();449    if (model == nullptr) {450        fprintf(stderr, "%s : failed to init model\n", __func__);451        return 1;452    }453 454    if (!llama_model_is_recurrent(model) && !llama_model_is_hybrid(model)) {455        fprintf(stderr, "%s : skipping for non-recurrent model\n", __func__);456        return 0;457    }458 459    for (uint8_t fill : { 0, 0x3e }) {460        fprintf(stderr, "%s : testing with cache fill 0x%02x\n", __func__, fill);461        if (test_rollback(params, model, fill) != 0) {462            return 1;463        }464    }465 466    return 0;467}468