Felipe97/llama-cpp-compiled

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