Felipe97/llama-cpp-compiled

1#include "llama-memory-hybrid.h"2 3#include "llama-impl.h"4#include "llama-model.h"5#include "llama-context.h"6 7//8// llama_memory_hybrid9//10 11llama_memory_hybrid::llama_memory_hybrid(12        const llama_model & model,13                            /* attn */14                ggml_type   type_k,15                ggml_type   type_v,16                     bool   v_trans,17                 uint32_t   kv_size,18                 uint32_t   n_pad,19                 uint32_t   n_swa,20           llama_swa_type   swa_type,21                            /* recurrent */22                ggml_type   type_r,23                ggml_type   type_s,24                 uint32_t   rs_size,25                            /* common */26                 uint32_t   n_seq_max,27                 uint32_t   n_rs_seq,28                     bool   offload,29                     bool   unified,30                            /* layer filters */31    const layer_filter_cb & filter_attn,32    const layer_filter_cb & filter_recr) :33    hparams(model.hparams),34    mem_attn(new llama_kv_cache(35        model,36        model.hparams,37        type_k,38        type_v,39        v_trans,40        offload,41        unified,42        kv_size,43        n_seq_max,44        n_pad,45        n_swa,46        swa_type,47        nullptr,48        filter_attn == nullptr ?49            [&](int32_t il) { return !hparams.is_recr(il); }50            : filter_attn,51        nullptr,52        nullptr53    )),54    mem_recr(new llama_memory_recurrent(55        model,56        type_r,57        type_s,58        offload,59        rs_size,60        n_seq_max,61        n_rs_seq,62        filter_recr == nullptr ?63            [&](int32_t il) { return hparams.is_recr(il); }64            : filter_recr65    )) {}66 67llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {68    do {69        balloc.split_reset();70 71        // follow the recurrent pattern for creating the ubatch splits72        std::vector<llama_ubatch> ubatches;73 74        while (true) {75            llama_ubatch ubatch;76 77            if (embd_all) {78                // if all tokens are output, split by sequence79                ubatch = balloc.split_seq(n_ubatch);80            } else {81                // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)82                const bool unified = (mem_attn->get_n_stream() == 1);83 84                // [TAG_RECURRENT_ROLLBACK_SPLITS]85                // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch86                //   so that the rollback snapshots remain valid87                const uint32_t n_rs_seq = mem_recr->n_rs_seq;88 89                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);90            }91 92            if (ubatch.n_tokens == 0) {93                break;94            }95 96            ubatches.push_back(std::move(ubatch)); // NOLINT97        }98 99        if (balloc.get_n_used() < balloc.get_n_tokens()) {100            // failed to find a suitable split101            break;102        }103 104        // prepare the recurrent batches first105        if (!mem_recr->prepare(ubatches)) {106            // TODO: will the recurrent cache be in an undefined context at this point?107            LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__);108            return std::make_unique<llama_memory_hybrid_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);109        }110 111        // prepare the attention cache112        auto heads_attn = mem_attn->prepare(ubatches);113        if (heads_attn.empty()) {114            LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__);115            return std::make_unique<llama_memory_hybrid_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);116        }117 118        return std::make_unique<llama_memory_hybrid_context>(119                this, std::move(heads_attn), std::move(ubatches));120    } while(false);121 122    return std::make_unique<llama_memory_hybrid_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);123}124 125llama_memory_context_ptr llama_memory_hybrid::init_full() {126    return std::make_unique<llama_memory_hybrid_context>(this);127}128 129llama_memory_context_ptr llama_memory_hybrid::init_update(llama_context * lctx, bool optimize) {130    return std::make_unique<llama_memory_hybrid_context>(this, lctx, optimize);131}132 133bool llama_memory_hybrid::get_can_shift() const {134    // Shifting is trivially supported for recurrent135    return mem_attn->get_can_shift();136}137 138void llama_memory_hybrid::clear(bool data) {139    mem_attn->clear(data);140    mem_recr->clear(data);141}142 143bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {144    // Try removing from the recurrent cache first since it may fail. If it does145    // fail, the cache will not have been mutated.146    if (!mem_recr->seq_rm(seq_id, p0, p1)) {147        return false;148    }149    return mem_attn->seq_rm(seq_id, p0, p1);150}151 152void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {153    mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1);154    mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1);155}156 157void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) {158    mem_attn->seq_keep(seq_id);159    mem_recr->seq_keep(seq_id);160}161 162void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {163    mem_attn->seq_add(seq_id, p0, p1, shift);164    mem_recr->seq_add(seq_id, p0, p1, shift);165}166 167void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {168    mem_attn->seq_div(seq_id, p0, p1, d);169    mem_recr->seq_div(seq_id, p0, p1, d);170}171 172llama_pos llama_memory_hybrid::seq_pos_min(llama_seq_id seq_id) const {173    // the min of the total cache is the max of the two caches' min values174    return std::max(mem_attn->seq_pos_min(seq_id), mem_recr->seq_pos_min(seq_id));175}176 177llama_pos llama_memory_hybrid::seq_pos_max(llama_seq_id seq_id) const {178    // the max of the total cache is the min of the two caches' max values179    return std::min(mem_attn->seq_pos_max(seq_id), mem_recr->seq_pos_max(seq_id));180}181 182std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid::memory_breakdown() const {183    std::map<ggml_backend_buffer_type_t, size_t> mb = mem_attn->memory_breakdown();184    for (const auto & buft_size : mem_recr->memory_breakdown()) {185        mb[buft_size.first] += buft_size.second;186    }187    return mb;188}189 190void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {191    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {192        mem_attn->state_write(io, seq_id, flags);193    }194    mem_recr->state_write(io, seq_id, flags);195}196 197void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {198    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {199        mem_attn->state_read(io, seq_id, flags);200    }201    mem_recr->state_read(io, seq_id, flags);202}203 204llama_kv_cache * llama_memory_hybrid::get_mem_attn() const {205    return mem_attn.get();206}207 208llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const {209    return mem_recr.get();210}211 212llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_status status) : status(status) {}213 214llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_hybrid * mem) :215    ctx_attn(mem->get_mem_attn()->init_full()),216    ctx_recr(mem->get_mem_recr()->init_full()),217    status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {218}219 220llama_memory_hybrid_context::llama_memory_hybrid_context(221        llama_memory_hybrid * mem,222              llama_context * lctx,223                       bool   optimize) :224    ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)),225    ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)),226    status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {227}228 229llama_memory_hybrid_context::llama_memory_hybrid_context(230              llama_memory_hybrid * mem,231                  slot_info_vec_t   sinfos_attn,232        std::vector<llama_ubatch>   ubatches) :233    ubatches(std::move(ubatches)),234    // note: here we copy the ubatches. not sure if this is ideal235    ctx_attn(new llama_kv_cache_context(mem->get_mem_attn(), std::move(sinfos_attn), this->ubatches)),236    ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)),237    status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) {238}239 240bool llama_memory_hybrid_context::next() {241    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);242 243    ctx_attn->next();244    ctx_recr->next();245 246    if (++i_next >= ubatches.size()) {247        return false;248    }249 250    return true;251}252 253bool llama_memory_hybrid_context::apply() {254    assert(!llama_memory_status_is_fail(status));255 256    bool res = true;257 258    res = res & ctx_attn->apply();259    res = res & ctx_recr->apply();260 261    return res;262}263 264llama_memory_status llama_memory_hybrid_context::get_status() const {265    return status;266}267 268const llama_ubatch & llama_memory_hybrid_context::get_ubatch() const {269    assert(status == LLAMA_MEMORY_STATUS_SUCCESS);270    return ubatches[i_next];271}272 273const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const {274    return static_cast<const llama_kv_cache_context *>(ctx_attn.get());275}276 277const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const {278    return static_cast<const llama_memory_recurrent_context *>(ctx_recr.get());279}280