Felipe97/llama-cpp-compiled

1#pragma once2 3#include "llama-batch.h"4#include "llama-graph.h"5#include "llama-memory.h"6 7#include <map>8#include <set>9#include <vector>10 11//12// llama_memory_recurrent13//14 15// TODO: extract the cache state used for graph computation into llama_memory_recurrent_context_i16//       see the implementation of llama_kv_cache_context_i for an example how to do it17class llama_memory_recurrent : public llama_memory_i {18public:19    llama_memory_recurrent(20            const llama_model & model,21                    ggml_type   type_r,22                    ggml_type   type_s,23                         bool   offload,24                     uint32_t   mem_size,25                     uint32_t   n_seq_max,26                     uint32_t   n_rs_seq,27        const layer_filter_cb & filter);28 29    ~llama_memory_recurrent() = default;30 31    //32    // llama_memory_i33    //34 35    llama_memory_context_ptr init_batch(36            llama_batch_allocr & balloc,37            uint32_t n_ubatch,38            bool embd_all) override;39 40    llama_memory_context_ptr init_full() override;41 42    llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;43 44    void clear(bool data) override;45 46    bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;47    void seq_cp  (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;48    void seq_keep(llama_seq_id seq_id)                                                          override;49    void seq_add (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, llama_pos shift) override;50    void seq_div (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, int d) override;51 52    llama_pos seq_pos_min(llama_seq_id seq_id) const override;53    llama_pos seq_pos_max(llama_seq_id seq_id) const override;54 55    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;56 57    bool prepare(const std::vector<llama_ubatch> & ubatches);58 59    // find a contiguous slot of memory cells and emplace the ubatch there60    bool find_slot(const llama_ubatch & ubatch);61 62    bool get_can_shift() const override;63 64    // state write/load65 66    void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;67    void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;68 69    uint32_t head = 0; // the location where the batch will be placed in the cache (see find_slot())70    uint32_t size = 0; // total number of cells, shared across all sequences71    uint32_t used = 0; // used cells (i.e. at least one seq_id)72 73    // number of recurrent-state snapshots per seq for rollback; tensors are widened to (1 + n_rs_seq) groups74    uint32_t n_rs_seq = 0;75 76    // per-seq rollback index77    std::vector<uint32_t> rs_idx;78 79    void set_rs_idx(llama_seq_id seq_id, uint32_t idx);80 81    // computed before each graph build82    uint32_t n = 0;83 84    // first zero-ed state85    int32_t rs_z = -1;86 87    // TODO: optimize for recurrent state needs88    struct mem_cell {89        llama_pos pos  = -1;90        int32_t   src  = -1; // used to know where states should be copied from91        int32_t   src0 = -1; // like src, but only used when setting the inputs (allowing to copy once)92        int32_t   tail = -1;93 94        std::set<llama_seq_id> seq_id;95 96        bool has_seq_id(const llama_seq_id & id) const {97            return seq_id.find(id) != seq_id.end();98        }99 100        bool is_empty() const {101            return seq_id.empty();102        }103 104        bool is_same_seq(const mem_cell & other) const {105            return seq_id == other.seq_id;106        }107    };108 109    std::vector<mem_cell> cells;110 111    // per layer112    std::vector<ggml_tensor *> r_l;113    std::vector<ggml_tensor *> s_l;114    // a second conv history that must stay replicated across devices, so it cannot share the r row115    std::vector<ggml_tensor *> p_l;116 117private:118    //const llama_model & model;119    const llama_hparams & hparams;120 121    const uint32_t n_seq_max = 1;122 123    // ggml contexts for the KV cache along with the allocated backend buffers:124    std::vector<std::pair<ggml_context_ptr, ggml_backend_buffer_ptr>> ctxs_bufs;125 126    size_t total_size() const;127 128    size_t size_r_bytes() const;129    size_t size_s_bytes() const;130    size_t size_p_bytes() const;131 132    void state_write_meta(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges, llama_seq_id seq_id = -1) const;133    void state_write_data(llama_io_write_i & io, const std::vector<std::pair<uint32_t, uint32_t>> & cell_ranges) const;134 135    bool state_read_meta(llama_io_read_i & io, uint32_t cell_count, llama_seq_id dest_seq_id = -1);136    bool state_read_data(llama_io_read_i & io, uint32_t cell_count);137};138 139class llama_memory_recurrent_context : public llama_memory_context_i {140public:141    // used for errors142    llama_memory_recurrent_context(llama_memory_status status);143 144    // used to create a full-cache or update context145    llama_memory_recurrent_context(146            llama_memory_recurrent * mem);147 148    // used to create a batch processing context from a batch149    llama_memory_recurrent_context(150            llama_memory_recurrent * mem,151            std::vector<llama_ubatch> ubatches);152 153    virtual ~llama_memory_recurrent_context();154 155    //156    // llama_memory_context_i157    //158 159    bool next()  override;160    bool apply() override;161 162    llama_memory_status  get_status() const override;163    const llama_ubatch & get_ubatch() const override;164 165    //166    // llama_memory_recurrent_context specific API167    //168 169    uint32_t get_n_rs() const;170    uint32_t get_head() const;171    int32_t  get_rs_z() const;172    uint32_t get_size() const;173 174    ggml_tensor * get_r_l(int32_t il) const;175    ggml_tensor * get_s_l(int32_t il) const;176    ggml_tensor * get_p_l(int32_t il) const;177 178    int32_t s_copy(int i) const;179 180private:181    const llama_memory_status status;182 183    llama_memory_recurrent * mem;184 185    size_t i_next = 0;186 187    std::vector<llama_ubatch> ubatches;188 189    //190    // data needed for building the compute graph for the current ubatch:191    // TODO: extract all the state like `head` and `n` here192    //193 194    const bool is_full = false;195};196