Felipe97/llama-cpp-compiled

1#pragma once2 3#include "llama-kv-cache.h"4 5#include <vector>6 7//8// llama_kv_cache_iswa9//10 11// utilizes two instances of llama_kv_cache12//   the first instance is for the non-SWA layers of the model and the second instance is for the SWA layers13 14class llama_kv_cache_iswa : public llama_memory_i {15public:16    llama_kv_cache_iswa(17            const llama_model & model,18                    ggml_type   type_k,19                    ggml_type   type_v,20                         bool   v_trans,21                         bool   offload,22                         bool   swa_full,23                         bool   unified,24                     uint32_t   kv_size,25                     uint32_t   n_seq_max,26                     uint32_t   n_ubatch,27                     uint32_t   n_pad,28               llama_memory_t   mem_other,29        const layer_filter_cb & filter,30        const  layer_reuse_cb & reuse,31        const  layer_share_cb & share);32 33    llama_kv_cache_iswa(34            const llama_model & model,35            const llama_hparams & hparams,36                    ggml_type   type_k,37                    ggml_type   type_v,38                         bool   v_trans,39                         bool   offload,40                         bool   swa_full,41                         bool   unified,42                     uint32_t   kv_size,43                     uint32_t   n_seq_max,44                     uint32_t   n_ubatch,45                     uint32_t   n_pad,46               llama_memory_t   mem_other,47        const layer_filter_cb & filter,48        const  layer_reuse_cb & reuse,49        const  layer_share_cb & share);50 51    ~llama_kv_cache_iswa() = default;52 53    //54    // llama_memory_i55    //56 57    llama_memory_context_ptr init_batch(58            llama_batch_allocr & balloc,59            uint32_t n_ubatch,60            bool embd_all) override;61 62    llama_memory_context_ptr init_full() override;63 64    llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override;65 66    bool get_can_shift() const override;67 68    void clear(bool data) override;69 70    bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;71    void seq_cp  (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override;72    void seq_keep(llama_seq_id seq_id)                                                          override;73    void seq_add (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, llama_pos shift) override;74    void seq_div (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1, int d) override;75 76    llama_pos seq_pos_min(llama_seq_id seq_id) const override;77    llama_pos seq_pos_max(llama_seq_id seq_id) const override;78 79    std::map<ggml_backend_buffer_type_t, size_t> memory_breakdown() const override;80 81    // state write/load82 83    void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override;84    void state_read (llama_io_read_i  & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override;85 86    //87    // llama_kv_cache_iswa specific API88    //89 90    llama_kv_cache * get_base() const;91    llama_kv_cache * get_swa () const;92 93private:94    const bool unified;95 96    std::unique_ptr<llama_kv_cache> kv_base;97    std::unique_ptr<llama_kv_cache> kv_swa;98};99 100class llama_kv_cache_iswa_context : public llama_memory_context_i {101public:102    using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;103 104    // used for errors105    llama_kv_cache_iswa_context(llama_memory_status status);106 107    // used to create a full-cache context108    llama_kv_cache_iswa_context(109            llama_kv_cache_iswa * kv);110 111    // used to create an update context112    llama_kv_cache_iswa_context(113            llama_kv_cache_iswa * kv,114            llama_context * lctx,115            bool optimize);116 117    // used to create a batch processing context from a batch118    llama_kv_cache_iswa_context(119            llama_kv_cache_iswa * kv,120            slot_info_vec_t sinfos_base,121            slot_info_vec_t sinfos_swa,122            std::vector<llama_ubatch> ubatches);123 124    virtual ~llama_kv_cache_iswa_context();125 126    //127    // llama_memory_context_i128    //129 130    bool next()  override;131    bool apply() override;132 133    llama_memory_status  get_status() const override;134    const llama_ubatch & get_ubatch() const override;135 136    //137    // llama_kv_cache_iswa_context specific API138    //139 140    const llama_kv_cache_context * get_base() const;141    const llama_kv_cache_context * get_swa()  const;142 143private:144    //llama_kv_cache_iswa * kv;145 146    // the index of the next ubatch to process147    size_t i_next = 0;148 149    std::vector<llama_ubatch> ubatches;150 151    const llama_memory_context_ptr ctx_base;152    const llama_memory_context_ptr ctx_swa;153 154    const llama_memory_status status;155};156