Felipe97/llama-cpp-compiled

1#pragma once2 3#include "llama-kv-cache.h"4 5#include <vector>6 7// llama_kv_cache_msa8 9// uses two instances of llama_kv_cache, one for K/V tensors, and one for the MSA indexer tensors10// both receive identical sequence operations and identical ubatches, so their cell layouts stay in synced.11// the context also exposes per-ubatch pos - cell translation maps populated from llama_kv_cells via12// llama_kv_cache::get_cells(), which the model graph uses to run MSA block selection in position space13 14class llama_kv_cache_msa : public llama_memory_i {15public:16    llama_kv_cache_msa(17            const llama_model & model,18                    ggml_type   type_k,19                    ggml_type   type_v,20                         bool   v_trans,21                         bool   offload,22                         bool   unified,23                     uint32_t   kv_size,24                     uint32_t   n_seq_max,25                     uint32_t   n_pad,26                     uint32_t   n_swa,27               llama_swa_type   swa_type,28        const layer_filter_cb & filter,29        const layer_filter_cb & filter_idx,30        const  layer_reuse_cb & reuse);31 32    ~llama_kv_cache_msa() = default;33 34    // llama_memory_i35 36    llama_memory_context_ptr init_batch(37            llama_batch_allocr & balloc,38            uint32_t n_ubatch,39            bool embd_all) override;40 41    llama_memory_context_ptr init_full() override;42 43    llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;44 45    bool get_can_shift() const override;46 47    void clear(bool data) override;48 49    bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;50    void seq_cp  (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;51    void seq_keep(llama_seq_id seq_id)                                                          override;52    void seq_add (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, llama_pos shift) override;53    void seq_div (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, int d) override;54 55    llama_pos seq_pos_min(llama_seq_id seq_id) const override;56    llama_pos seq_pos_max(llama_seq_id seq_id) const override;57 58    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;59 60    // state write/load61 62    void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;63    void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;64 65    // llama_kv_cache_msa specific API66 67    llama_kv_cache * get_base() const;68    llama_kv_cache * get_idx () const;69 70    uint32_t       get_n_pad()    const { return n_pad; }71    uint32_t       get_n_seq_max() const { return n_seq_max; }72    uint32_t       get_n_swa()    const { return n_swa; }73    llama_swa_type get_swa_type() const { return swa_type; }74 75private:76    // keep the indexer KV cache hparams instance here as llama_kv_cache stores only a reference77    llama_hparams hparams_idx;78 79    const uint32_t n_stream  = 1;80    const uint32_t n_seq_max = 1;81    const uint32_t n_pad     = 1;82 83    const uint32_t       n_swa    = 0;84    const llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE;85 86    std::unique_ptr<llama_kv_cache> kv_base;87    std::unique_ptr<llama_kv_cache> kv_idx;88};89 90class llama_kv_cache_msa_context : public llama_memory_context_i {91public:92    using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;93 94    // used for errors95    llama_kv_cache_msa_context(llama_memory_status status);96 97    // used to create a full-cache context98    llama_kv_cache_msa_context(99            llama_kv_cache_msa * kv);100 101    // used to create an update context102    llama_kv_cache_msa_context(103            llama_kv_cache_msa * kv,104            llama_context * lctx,105            bool optimize);106 107    // used to create a batch processing context from a batch108    llama_kv_cache_msa_context(109            llama_kv_cache_msa * kv,110            slot_info_vec_t sinfos_base,111            slot_info_vec_t sinfos_idx,112            std::vector<llama_ubatch> ubatches);113 114    virtual ~llama_kv_cache_msa_context();115 116    // llama_memory_context_i117 118    bool next()  override;119    bool apply() override;120 121    llama_memory_status  get_status() const override;122    const llama_ubatch & get_ubatch() const override;123 124    // llama_kv_cache_msa_context specific API125 126    const llama_kv_cache_context * get_base() const;127    const llama_kv_cache_context * get_idx () const;128 129    // max position currently present in the cache plus one, padded MSA blocks are defined over token positions130    // so the block-selection tensors are sized by this value rather than by the number of cells131    uint32_t get_n_pos() const;132 133    // position <-> cell translation maps, populated from the base cache cells134    // the model graph relates cache contents to token positions only through these per ubatch inputs135    // value for empty or other-sequence cells is 0 so consumers must mask them136    void set_input_cell_pos(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const;137    // positions without a cell map to cell 0, consumers must mask them assumes one sequence per stream138    void set_input_pos_slot(ggml_tensor * dst, const llama_ubatch * ubatch) const;139    void set_input_pos_mask(ggml_tensor * dst, const llama_ubatch * ubatch) const;140 141private:142    llama_kv_cache_msa * kv;143 144    // the index of the next ubatch to process145    size_t i_next = 0;146 147    std::vector<llama_ubatch> ubatches;148 149    const llama_memory_context_ptr ctx_base;150    const llama_memory_context_ptr ctx_idx;151 152    const llama_memory_status status;153};154