Felipe97/llama-cpp-compiled

1#pragma once2 3#include "llama.h"4#include "llama-cparams.h"5 6#include <bitset>7#include <cassert>8#include <cstring>9#include <limits>10#include <set>11#include <vector>12 13struct llama_kv_cell_ext {14    // 2D spatial positions, typically used for M-RoPE15    llama_pos x = 0;16    llama_pos y = 0;17 18    // when tok = LLAMA_TOKEN_NULL when the cell is produced by embedding input (i.e. multimodal)19    // use case: n-gram embeddings hash20    llama_token tok = LLAMA_TOKEN_NULL;21 22    // return true if the current 2D spatial position is greater than other23    bool is_2d_gt(llama_pos ox, llama_pos oy) const {24        return (y > oy) || (y == oy && x > ox);25    }26 27    void reset() {28        static_assert(std::is_trivially_copyable_v<llama_kv_cell_ext>);29 30        *this = llama_kv_cell_ext{};31    }32};33 34// meta information about KV cells that can be part of multiple sequences at the same time35// TODO: add unit tests36class llama_kv_cells {37public:38    using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;39 40    void reset() {41        for (uint32_t i = 0; i < pos.size(); ++i) {42            pos[i]   = -1;43            ext[i].reset();44            shift[i] =  0;45            seq[i].reset();46        }47 48        has_shift = false;49 50        used.clear();51 52        for (uint32_t s = 0; s < LLAMA_MAX_SEQ; ++s) {53            seq_pos[s].clear();54        }55    }56 57    void reset_shift() {58        has_shift = false;59 60        for (uint32_t i = 0; i < shift.size(); ++i) {61            shift[i] = 0;62        }63    }64 65    uint32_t size() const {66        return pos.size();67    }68 69    void resize(uint32_t n) {70        pos.resize(n);71        ext.resize(n);72        shift.resize(n);73        seq.resize(n);74 75        reset();76    }77 78    bool is_empty(uint32_t i) const {79        assert(i < pos.size());80        assert((pos[i] < 0 && pos[i] == -1) || pos[i] >= 0);81 82        return pos[i] == -1;83    }84 85    uint32_t get_used() const {86        return used.size();87    }88 89    // the index of the first cell that is used90    // return 0 if no cells are used91    uint32_t used_min() const {92        return used.empty() ? 0 : *used.begin();93    }94 95    // the index of the last cell that is used + 196    // return 0 if no cells are used97    uint32_t used_max_p1() const {98        return used.empty() ? 0 : *used.rbegin() + 1;99    }100 101    bool get_has_shift() const {102        return has_shift;103    }104 105    // move cell isrc to idst (used during defrag)106    //void mv(uint32_t isrc, uint32_t idst) {107    //    assert(isrc < pos.size());108    //    assert(idst < pos.size());109 110    //    assert(pos[idst] == -1);111    //    assert(pos[isrc] != -1);112 113    //    pos  [idst] = pos  [isrc];114    //    shift[idst] = shift[isrc];115    //    seq  [idst] = seq  [isrc];116 117    //    pos  [isrc] = -1;118    //    shift[isrc] =  0;119    //    seq  [isrc].reset();120 121    //    used.erase (isrc);122    //    used.insert(idst);123    //}124 125    // copy the state of cells [i, i + n) (used for save/restore the state of the cells)126    llama_kv_cells cp(uint32_t i, uint32_t n) const {127        assert(i + n <= pos.size());128 129        llama_kv_cells res;130 131        res.resize(n);132 133        for (uint32_t j = 0; j < n; ++j) {134            const auto idx = i + j;135 136            res.pos[j] = pos[idx];137            res.ext[j] = ext[idx];138            res.seq[j] = seq[idx];139 140            assert(shift[idx] == 0);141        }142 143        return res;144    }145 146    // copy the state of cells [idxs[0], idxs[1], ..., idxs[idxs.size() - 1])147    llama_kv_cells cp(const std::vector<uint32_t> & idxs) const {148        llama_kv_cells res;149 150        res.resize(idxs.size());151 152        for (uint32_t j = 0; j < idxs.size(); ++j) {153            const auto idx = idxs[j];154 155            res.pos[j] = pos[idx];156            res.ext[j] = ext[idx];157            res.seq[j] = seq[idx];158 159            assert(shift[idx] == 0);160        }161 162        return res;163    }164 165    // set the state of cells [i, i + other.pos.size()) (used for save/restore the state of the cells)166    void set(uint32_t i, const llama_kv_cells & other) {167        assert(i + other.pos.size() <= pos.size());168 169        for (uint32_t j = 0; j < other.pos.size(); ++j) {170            const auto idx = i + j;171 172            if (pos[idx] == -1 && other.pos[j] != -1) {173                used.insert(i + j);174            }175 176            if (pos[idx] != -1 && other.pos[j] == -1) {177                used.erase(i + j);178            }179 180            if (pos[idx] != -1) {181                seq_pos_rm(i + j);182            }183 184            pos[idx] = other.pos[j];185            ext[idx] = other.ext[j];186            seq[idx] = other.seq[j];187 188            if (pos[idx] != -1) {189                seq_pos_add(i + j);190            }191 192            assert(shift[idx] == 0);193        }194    }195 196    // set the state of cells [idxs[0], idxs[1], ..., idxs[idxs.size() - 1])197    void set(const std::vector<uint32_t> & idxs, const llama_kv_cells & other) {198        assert(idxs.size() == other.pos.size());199 200        for (uint32_t j = 0; j < other.pos.size(); ++j) {201            const auto idx = idxs[j];202 203            if (pos[idx] == -1 && other.pos[j] != -1) {204                used.insert(idx);205            }206 207            if (pos[idx] != -1 && other.pos[j] == -1) {208                used.erase(idx);209            }210 211            if (pos[idx] != -1) {212                seq_pos_rm(idx);213            }214 215            pos[idx] = other.pos[j];216            ext[idx] = other.ext[j];217            seq[idx] = other.seq[j];218 219            if (pos[idx] != -1) {220                seq_pos_add(idx);221            }222 223            assert(shift[idx] == 0);224        }225    }226 227    // clear a non-empty cell228    void rm(uint32_t i) {229        assert(i < pos.size());230        assert(pos[i] != -1);231 232        seq_pos_rm(i);233        seq[i].reset();234 235        pos[i] = -1;236        ext[i].reset();237        shift[i] = 0;238 239        used.erase(i);240    }241 242    // note: call only if the cell has seq_id243    // return true if the cell becomes empty244    bool seq_rm(uint32_t i, llama_seq_id seq_id) {245        assert(i < pos.size());246        assert(seq[i].test(seq_id));247        assert(pos[i] != -1);248        assert(seq_id >= 0);249 250        seq[i].reset(seq_id);251        seq_pos_dec(seq_id, i);252 253        if (seq[i].none()) {254            pos[i] = -1;255            ext[i].reset();256            shift[i] = 0;257 258            used.erase(i);259 260            return true;261        }262 263        return false;264    }265 266    // return true if the cell becomes empty (i.e. it did not contain seq_id before the call)267    bool seq_keep(uint32_t i, llama_seq_id seq_id) {268        assert(i < pos.size());269 270        if (seq[i].test(seq_id)) {271            seq_pos_rm(i);272            seq[i].reset();273 274            seq[i].set(seq_id);275            seq_pos_inc(seq_id, i);276 277            return false;278        }279 280        if (seq[i].any()) {281            seq_pos_rm(i);282            seq[i].reset();283 284            pos[i] = -1;285            ext[i].reset();286            shift[i] = 0;287 288            used.erase(i);289 290            return true;291        }292 293        assert(pos[i] == -1);294 295        return false;296    }297 298    // number of different sequences in the cell299    int seq_count(uint32_t i) const {300        assert(i < pos.size());301        assert(pos[i] != -1);302 303        return seq[i].count();304    }305 306    // the full set of sequences this cell is visible to307    const seq_set_t & seq_get_all(uint32_t i) const {308        assert(i < pos.size());309 310        return seq[i];311    }312 313    // check if the cell contains seq_id314    bool seq_has(uint32_t i, llama_seq_id seq_id) const {315        assert(i < pos.size());316        assert(seq_id >= 0);317 318        return seq[i].test(seq_id);319    }320 321    // the token of the cell of sequence seq_id at the largest position <= p322    // when several cells share that position, the one with the highest index wins323    // return LLAMA_TOKEN_NULL if the sequence has no cell at or before p324    // note: used by n-gram input embeddings to recover the tokens preceding a ubatch325    llama_token seq_pos_tok_le(llama_seq_id seq_id, llama_pos p) const {326        assert(seq_id >= 0);327        assert(seq_id < LLAMA_MAX_SEQ);328 329        const auto & sp = seq_pos[seq_id];330 331        auto it = sp.upper_bound({ p, std::numeric_limits<uint32_t>::max() });332        if (it == sp.begin()) {333            return LLAMA_TOKEN_NULL;334        }335 336        return ext[(--it)->second].tok;337    }338 339    // note: call only if the cell is not empty and the seq_id is not in the cell340    void seq_add(uint32_t i, llama_seq_id seq_id) {341        assert(i < pos.size());342        assert(pos[i] != -1);343        assert(!seq[i].test(seq_id));344 345        seq[i].set(seq_id);346        seq_pos_inc(seq_id, i);347    }348 349    // return the sequence id of this cell350    // note: call only for cells with exactly one sequence351    llama_seq_id seq_get(uint32_t i) const {352        assert(seq[i].count() == 1);353 354        for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {355            if (seq[i].test(s)) {356                return s;357            }358        }359 360        return -1;361    }362 363    // the minimum position of sequence seq_id currently present in any of the cells364    // return -1 if the sequence is not present365    llama_pos seq_pos_min(llama_seq_id seq_id) const {366        assert(seq_id >= 0);367        assert(seq_id < LLAMA_MAX_SEQ);368 369        if (seq_pos[seq_id].empty()) {370            return -1;371        }372 373        return seq_pos[seq_id].begin()->first;374    }375 376    // the maximum position of sequence seq_id currently present in any of the cells377    // return -1 if the sequence is not present378    llama_pos seq_pos_max(llama_seq_id seq_id) const {379        assert(seq_id >= 0);380        assert(seq_id < LLAMA_MAX_SEQ);381 382        if (seq_pos[seq_id].empty()) {383            return -1;384        }385 386        return seq_pos[seq_id].rbegin()->first;387    }388 389    // note: call only if the cell is not empty390    llama_pos pos_get(uint32_t i) const {391        assert(i < pos.size());392        assert(pos[i] != -1);393 394        return pos[i];395    }396 397    const llama_kv_cell_ext & ext_get(uint32_t i) const {398        assert(i < pos.size());399        assert(pos[i] != -1);400 401        return ext[i];402    }403 404    // note: call only if the cell is not empty405    llama_pos get_shift(uint32_t i) const {406        assert(i < pos.size());407        assert(pos[i] != -1);408 409        return shift[i];410    }411 412    // check if a cell is not empty and its position is within [p0, p1)413    bool pos_in(uint32_t i, llama_pos p0, llama_pos p1) const {414        assert(i < pos.size());415 416        return pos[i] >= p0 && pos[i] < p1;417    }418 419    // set the position of an empty cell420    // does not modify "has_shift"421    // note: call only if the cell is empty422    void pos_set(uint32_t i, llama_pos p) {423        assert(i < pos.size());424        assert(pos[i] == -1);425        assert(seq[i].none());426 427        pos[i] = p;428 429        used.insert(i);430    }431 432    void ext_set(uint32_t i, llama_kv_cell_ext p) {433        assert(i < ext.size());434        ext[i] = p;435    }436 437    // pos[i] = pos[i] + d438    // sets "has_shift" to true439    // note: call only if the cell is not empty440    bool pos_add(uint32_t i, llama_pos d) {441        assert(i < pos.size());442        assert(pos[i] != -1);443 444        seq_pos_rm(i);445 446        pos[i]   += d;447        shift[i] += d;448 449        has_shift = true;450 451        if (pos[i] < 0) {452            seq[i].reset();453            pos[i] = -1;454            shift[i] = 0;455 456            used.erase(i);457 458            return true;459        }460 461        seq_pos_add(i);462 463        return false;464    }465 466    // pos[i] = pos[i] / d467    // sets "has_shift" to true468    // note: call only if the cell is not empty469    void pos_div(uint32_t i, int d) {470        assert(i < pos.size());471        assert(pos[i] != -1);472 473        const llama_pos p_old = pos[i];474 475        seq_pos_rm(i);476 477        pos[i]   /= d;478        shift[i] += p_old - pos[i];479 480        seq_pos_add(i);481 482        has_shift = true;483    }484 485private:486    bool has_shift = false;487 488    // set of indices of used cells (i.e. pos[i] != -1, allowed to not have any seq_id)489    std::set<uint32_t> used;490 491    std::vector<llama_pos> pos;492 493    // stores extra info per cell494    std::vector<llama_kv_cell_ext> ext;495 496    // this array accumulates any applied shifts to the pos array since the last reset_shift() call497    // this is used to queue multiple updates to the pos array, which in the end can be applied in one go:498    //499    //   cells.pos_add(x, shift_x);500    //   cells.pos_div(y, shift_y);501    //   ...502    //503    //   if (cells.has_shift()) {504    //      for (int i = 0; i < n; ++i) {505    //          auto shift_i = cells.get_shift(i);506    //          ...507    //      }508    //      cells.reset_shift();509    //   }510    //511    std::vector<llama_pos> shift;512 513    // the bitset seq[i] tells us which sequences are currently occupying the i-th cell514    std::vector<seq_set_t> seq;515 516    // the set seq_pos[s] holds one (pos, cell) pair per cell that carries sequence s, ordered by position517    // this way seq_pos[s].begin() and seq_pos[s].rbegin() give us the min/max positions currently in the cache518    // and upper_bound() on a position finds the nearest cell of the sequence in logarithmic time519    //520    // the cell index is part of the key because a position can occur more than once for the same seq:521    //  - during performing a cache reuse via (rm + add)522    //  - some vision models have input embeddings with repeating positions523    //524    std::set<std::pair<llama_pos, uint32_t>> seq_pos[LLAMA_MAX_SEQ];525 526    // helper functions for updating `seq_pos`, once cell at a time:527 528    void seq_pos_dec(llama_seq_id s, uint32_t i) {529        const auto n = seq_pos[s].erase({ pos[i], i });530        assert(n == 1);531        GGML_UNUSED(n);532    }533 534    void seq_pos_inc(llama_seq_id s, uint32_t i) {535        seq_pos[s].insert({ pos[i], i });536    }537 538    // remove cell i539    void seq_pos_rm(uint32_t i) {540        for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {541            if (seq[i].test(s)) {542                seq_pos_dec(s, i);543            }544        }545    }546 547    // add cell i548    void seq_pos_add(uint32_t i) {549        for (int s = 0; s < LLAMA_MAX_SEQ; ++s) {550            if (seq[i].test(s)) {551                seq_pos_inc(s, i);552            }553        }554    }555};556 557using llama_kv_cells_vec = std::vector<llama_kv_cells>;558