Felipe97/llama-cpp-compiled

1#pragma once2 3#include "llama.h"4 5#include "llama-cparams.h"6 7#include <array>8#include <vector>9#include <set>10#include <bitset>11#include <memory>12#include <unordered_map>13 14// keep this struct lightweight15struct llama_ubatch {16    bool equal_seqs() const {17        return b_equal_seqs != 0;18    }19 20    // typical for M-RoPE cases:21    //   0 - sequential position of the tokens/embeddings in the sequence22    //   1 - y position in the image23    //   2 - x position in the image24    //   3 - other25    bool is_pos_2d() const {26        // TODO @ngxson : we may need to check for model arch when more models use >1 positions27        return n_pos >= 3;28    }29 30    uint32_t b_equal_seqs; // note: this is a boolean, but we use an int32_t for alignment31                           //       otherwise address sanitizer complains32    // TODO: whole_seqs for embeddings?33 34    uint32_t n_tokens;     // total tokens (n_seq_tokens * n_seqs)35    uint32_t n_seq_tokens; // tokens per sequence set36    uint32_t n_seqs;       // sequence sets in the ubatch37    uint32_t n_seqs_unq;   // unique sequence ids in the ubatch38    uint32_t n_pos;        // number of position inputs for each token/embedding39 40    // seq_id_unq: unique sequence ids in the ubatch41    // seq_idx:    indices of the unique sequence ids in the ubatch in [0, n_seqs_unq)42    //             used for extracting sequence pooled embeddings43 44    //                          // size               | idx | val45    llama_token  *  token;      // [n_tokens]         | i   | id, token46    float        *  embd;       // [n_embd, n_tokens] | i   | embd47    llama_pos    *  pos;        // [n_tokens*n_pos]   | i   | pos48    int32_t      *  n_seq_id;   // [n_tokens]         | i   | -49    llama_seq_id ** seq_id;     // [n_tokens]         | s   | s0, s1, seq_id50    llama_seq_id *  seq_id_unq; // [n_seqs_unq]       | s   | seq_id51    int32_t      *  seq_idx;    // [LLAMA_MAX_SEQ]    | -   | seq_idx52    int8_t       *  output;     // [n_tokens]         | i   | -53 54    struct data_t {55        std::vector<llama_token>    token;56        std::vector<float>          embd;57        std::vector<llama_pos>      pos;58        std::vector<int32_t>        n_seq_id;59        std::vector<llama_seq_id *> seq_id;      // these point into the seq_id_data below60        std::vector<llama_seq_id>   seq_id_unq;61        std::vector<int32_t>        seq_idx;62        std::vector<int8_t>         output;63 64        std::vector<llama_seq_id> seq_id_data;65    };66 67    // the llama_ubatch pointers above point to this data if set. otherwise - point to external non-owning data68    std::shared_ptr<data_t> data;69};70 71// a helper for sanitizing, fulfilling and splitting a batch72class llama_batch_allocr {73public:74    llama_batch_allocr(uint32_t n_pos_per_embd);75 76    // sanitize and auto-gen missing data in the input batch77    // memory is optional. if provided will be used to check for sequence continuity and to determine the positions78    bool init(79            const llama_batch & batch_inp,80            const llama_vocab & vocab,81            const llama_memory_i * memory,82            uint32_t n_embd,83            uint32_t n_seq_max,84            bool output_all);85 86    const llama_batch & get_batch() const;87 88    uint32_t get_n_tokens()  const;89    uint32_t get_n_outputs() const;90    uint32_t get_n_used()    const;91 92    // the array of output indices in the order they were encountered during the ubatch splitting93    std::vector<int32_t> & get_out_ids();94 95    // min/max positions of each sequence in the current ubatch96    llama_pos seq_pos_min(llama_seq_id seq_id) const;97    llama_pos seq_pos_max(llama_seq_id seq_id) const;98 99    // call once before splitting the batch to reset the internal state100    void split_reset();101 102    // simple split, unknown number of sequence sets of unequal lengths103    llama_ubatch split_simple(uint32_t n_ubatch);104 105    // make ubatches of equal-length sequences sets106    // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids107    // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch108    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);109 110    // sequence-set-wise split - each ubatch contains a single sequence-set111    llama_ubatch split_seq(uint32_t n_ubatch);112 113    // a helper method for creating a well-defined ubatch of tokens114    // TODO: support embeddings if needed in the future115    llama_ubatch ubatch_reserve(uint32_t n_seq_tokens, uint32_t n_seqs);116 117private:118    void clear();119 120    // create the next ubatch based on the provided batch indices (idxs) and the number of sequence sets (n_seqs)121    // return llama_ubatch.n_tokens == 0 if the entire batch was consumed122    llama_ubatch ubatch_add(const std::vector<int32_t> & idxs, uint32_t n_seqs, bool equal_seqs);123 124    // for debugging, start with LLAMA_BATCH_DEBUG=2125    void ubatch_print(const llama_ubatch & ubatch, int debug);126 127    llama_batch batch;128 129    // only for debugging purposes130    const llama_vocab * vocab;131 132    // TODO: this is more of a temporary solution until we have a better way to handle multiple positions per token/embd133    //       ref: https://github.com/ggml-org/llama.cpp/issues/13694#issuecomment-2983871762134    const uint32_t n_pos_per_embd;135 136    uint32_t n_embd;137    uint32_t n_seq_max;138    uint32_t n_outputs;139 140    std::array<llama_seq_id, 1> seq_id_0 = {{ 0 }}; // default sequence id141 142    std::vector<llama_pos>      pos;143    std::vector<int32_t>        n_seq_id;144    std::vector<llama_seq_id *> seq_id;145    std::vector<llama_seq_id>   seq_id_unq;146    std::vector<int32_t>        seq_idx;147    std::vector<int8_t>         output;148 149    using pos_set_t = std::set<llama_pos>;150    using seq_cpl_t = std::vector<bool>;151 152    // helper flag to quickly determine if there are any coupled sequences in the batch153    bool has_cpl = false;154 155    std::vector<pos_set_t> seq_pos; // seq_pos[s]: the set of positions in sequence s156    std::vector<seq_cpl_t> seq_cpl; // seq_cpl[s0][s1]: if sequence s0 is coupled to sequence s1157 158    using idx_vec_t = std::vector<int32_t>;159    using seq_set_t = std::bitset<LLAMA_MAX_SEQ>;160 161    std::vector<seq_set_t> seq_set; // seq_set[i]: the sequence set of token i162 163    std::unordered_map<seq_set_t, idx_vec_t> seq_set_map; // the indices at which the sequence set appears164 165    // batch indices of the output166    std::vector<int32_t> out_ids;167 168    uint32_t n_used;169 170    // used[i] indicates if token i has already been used in a previous ubatch171    std::vector<bool> used;172 173    int debug;174};175