Felipe97/llama-cpp-compiled

1#include "llama-memory-hybrid-idx.h"2 3#include "llama-impl.h"4#include "llama-batch.h"5#include "llama-io.h"6#include "llama-model.h"7 8 9#include <algorithm>10#include <cassert>11#include <cmath>12#include <iterator>13#include <stdexcept>14 15//16// llama_memory_hybrid_idx17//18 19llama_memory_hybrid_idx::llama_memory_hybrid_idx(20        const llama_model & model,21                            /* attn */22                ggml_type   type_k,23                ggml_type   type_v,24                     bool   v_trans,25                 uint32_t   kv_size,26                 uint32_t   n_pad,27                 uint32_t   n_swa,28           llama_swa_type   swa_type,29                            /* recurrent */30                ggml_type   type_r,31                ggml_type   type_s,32                 uint32_t   rs_size,33                            /* common */34                 uint32_t   n_seq_max,35                 uint32_t   n_rs_seq,36                     bool   offload,37                     bool   unified,38                            /* layer filters */39    const layer_filter_cb & filter_attn,40    const layer_filter_cb & filter_recr,41    const layer_filter_cb & filter_idx) :42    llama_memory_hybrid(43        model,44        type_k, type_v, v_trans, kv_size, n_pad, n_swa, swa_type,45        type_r, type_s, rs_size,46        n_seq_max, n_rs_seq, offload, unified,47        filter_attn, filter_recr),48    hparams_idx(model.hparams),49    mem_idx(filter_idx == nullptr ? nullptr : [&] {50        // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own51        std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);52        hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;53 54        // the cached indexer keys are raw, rotation happens after pooling at read time, so a55        // K-shift must not rotate them while the stream copies in the same update still apply56        hparams_idx.rope_type = LLAMA_ROPE_TYPE_NONE;57 58        // fool llama_kv_cache into thinking this is a MLA cache, so it won't cache V tensors59        hparams_idx.n_embd_head_k_mla_impl = model.hparams.indexer_head_size;60        hparams_idx.n_embd_head_v_mla_impl = model.hparams.indexer_head_size;61 62        LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);63 64        return new llama_kv_cache(65            model, hparams_idx, type_k, type_v, v_trans, offload, unified,66            kv_size, n_seq_max, n_pad, n_swa, swa_type,67            nullptr, filter_idx, nullptr, nullptr, "idx_");68    }()) {}69 70llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {71    // note: repeats llama_memory_hybrid::init_batch, as the indexer needs the attention slot infos that the base context hides72    do {73        balloc.split_reset();74 75        // follow the recurrent pattern for creating the ubatch splits76        std::vector<llama_ubatch> ubatches;77 78        while (true) {79            llama_ubatch ubatch;80 81            if (embd_all) {82                // if all tokens are output, split by sequence83                ubatch = balloc.split_seq(n_ubatch);84            } else {85                // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice)86                const bool unified = (get_mem_attn()->get_n_stream() == 1);87 88                // [TAG_RECURRENT_ROLLBACK_SPLITS]89                // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch90                //   so that the rollback snapshots remain valid91                const uint32_t n_rs_seq = get_mem_recr()->n_rs_seq;92 93                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);94            }95 96            if (ubatch.n_tokens == 0) {97                break;98            }99 100            ubatches.push_back(std::move(ubatch)); // NOLINT101        }102 103        if (balloc.get_n_used() < balloc.get_n_tokens()) {104            // failed to find a suitable split105            break;106        }107 108        // prepare the recurrent batches first109        if (!get_mem_recr()->prepare(ubatches)) {110            // TODO: will the recurrent cache be in an undefined context at this point?111            LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__);112            return std::make_unique<llama_memory_hybrid_idx_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);113        }114 115        // prepare the attention cache116        auto heads_attn = get_mem_attn()->prepare(ubatches);117        if (heads_attn.empty()) {118            LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__);119            return std::make_unique<llama_memory_hybrid_idx_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);120        }121 122        // the indexer uses the attention cache's slot layout; a separate one can drift from it123        llama_kv_cache::slot_info_vec_t heads_idx;124        if (mem_idx) {125            heads_idx = heads_attn;126        }127 128        return std::make_unique<llama_memory_hybrid_idx_context>(129                this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches));130    } while(false);131 132    return std::make_unique<llama_memory_hybrid_idx_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);133}134 135llama_memory_context_ptr llama_memory_hybrid_idx::init_full() {136    return std::make_unique<llama_memory_hybrid_idx_context>(this);137}138 139llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lctx, bool optimize) {140    return std::make_unique<llama_memory_hybrid_idx_context>(this, lctx, optimize);141}142 143void llama_memory_hybrid_idx::clear(bool data) {144    llama_memory_hybrid::clear(data);145 146    if (mem_idx) {147        mem_idx->clear(data);148    }149}150 151bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {152    // same order as llama_memory_hybrid::seq_rm: the recurrent cache can refuse, so try it first153    if (!get_mem_recr()->seq_rm(seq_id, p0, p1)) {154        return false;155    }156 157    if (mem_idx) {158        mem_idx->seq_rm(seq_id, p0, p1);159    }160 161    return get_mem_attn()->seq_rm(seq_id, p0, p1);162}163 164void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {165    llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1);166 167    if (mem_idx) {168        mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);169    }170}171 172void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {173    llama_memory_hybrid::seq_keep(seq_id);174 175    if (mem_idx) {176        mem_idx->seq_keep(seq_id);177    }178}179 180void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {181    llama_memory_hybrid::seq_add(seq_id, p0, p1, shift);182 183    if (mem_idx) {184        mem_idx->seq_add(seq_id, p0, p1, shift);185    }186}187 188void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {189    llama_memory_hybrid::seq_div(seq_id, p0, p1, d);190 191    if (mem_idx) {192        mem_idx->seq_div(seq_id, p0, p1, d);193    }194}195 196std::map<ggml_backend_buffer_type_t, size_t> llama_memory_hybrid_idx::memory_breakdown() const {197    std::map<ggml_backend_buffer_type_t, size_t> mb = llama_memory_hybrid::memory_breakdown();198 199    if (mem_idx) {200        for (const auto & buft_size : mem_idx->memory_breakdown()) {201            mb[buft_size.first] += buft_size.second;202        }203    }204 205    return mb;206}207 208void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const {209    llama_memory_hybrid::state_write(io, seq_id, flags);210 211    // [TAG_HYBRID_IDX_STATE] the indexer section goes last, so it is a pure suffix: an old reader stops early instead of misparsing it212    // The indexer mirrors the attention cache, so it uses the same PARTIAL_ONLY gate.213    if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {214        if (mem_idx) {215            mem_idx->state_write(io, seq_id, flags);216        }217    }218 219}220 221void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {222    // note: repeats llama_memory_hybrid::state_read223    // the indexer needs the attention cache's cells, and a half-failed restore must leave all three caches alike224 225    // [TAG_HYBRID_IDX_SINFO]226    // the indexer restore adopts the attention cache's layout instead of searching for cells of its own227    // two find_slot calls agree only while both caches see the same occupancy, which a restore cannot promise228    llama_kv_cache::slot_info_vec_t sinfos_attn;229 230    try {231        if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {232            get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr);233        }234 235        get_mem_recr()->state_read(io, seq_id, flags);236 237        // [TAG_HYBRID_IDX_STATE] must mirror the write order in state_write238        if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {239            if (mem_idx) {240                mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);241            }242        }243 244    } catch (...) {245        // a half-restored context is the one state the indexer cannot fix by itself: attention holds new cells, the indexer old ones246        // drop what was being restored from all of them, which is a state they do agree on.247        state_drop(seq_id);248 249        throw;250    }251}252 253void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) {254    // dropped directly, not via seq_rm: the recurrent cache may refuse it and then only the other two get cleared255    if (seq_id < 0) {256        clear(true);257 258        return;259    }260 261    get_mem_attn()->seq_rm(seq_id, -1, -1);262    get_mem_recr()->seq_rm(seq_id, -1, -1);263 264    if (mem_idx) {265        mem_idx->seq_rm(seq_id, -1, -1);266    }267}268 269llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const {270    return mem_idx.get();271}272 273void llama_memory_hybrid_idx::set_input_qsa(274        ggml_tensor * cell_blk,275        ggml_tensor * blk_cells,276        ggml_tensor * blk_pos,277        ggml_tensor * bias,278        const llama_ubatch * ubatch,279        uint32_t ratio,280        bool blk_bias) const {281    GGML_ASSERT(ratio > 0);282    GGML_ASSERT(get_mem_idx() != nullptr);283 284    GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer));285 286    const int64_t n_kv     = cell_blk->ne[0];287    const int64_t n_ns     = cell_blk->ne[1];        // streams in this ubatch288    const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns);289    const int64_t n_tokens = ubatch->n_tokens;290    const int64_t r        = ratio;291 292    GGML_ASSERT(n_tokens % n_ns == 0);293    const int64_t n_tps = n_tokens/n_ns;             // tokens per stream294 295    int32_t * dst_cell_blk  = (int32_t *) cell_blk->data;296    int32_t * dst_blk_cells = (int32_t *) blk_cells->data;297    int32_t * dst_blk_pos   = (int32_t *) blk_pos->data;298    float   * dst_bias      = (float   *) bias->data;299 300    // a block is keyed on (sequence set, index bucket): a unified cache counts every sequence301    // from zero, so the bucket alone would pool two sequences into one block302    GGML_ASSERT(r <= 64);303    const uint64_t slots_full = r == 64 ? ~uint64_t(0) : ((uint64_t(1) << r) - 1);304 305    // TODO: this runs per ubatch and is O(n_kv) per stream, about 865 us at 33k context. the cost306    //       is the per-cell scan rather than these allocations, so hoisting them buys nothing307    std::vector<int32_t>  blk_of(n_kv);308    std::vector<int32_t>  cell_grp(n_kv);309    std::vector<int32_t>  grp_head(n_blocks);310    std::vector<int32_t>  grp_next;311    std::vector<int32_t>  grp_first;312    std::vector<int32_t>  grp_slot0;313    std::vector<uint64_t> grp_slots;314    std::vector<int32_t>  grp_bid;315    std::vector<int32_t>  bid_idx;316    std::vector<int32_t>  bid_cell;317    std::vector<int32_t>  bid_slot0;318 319    std::vector<int32_t> order;320    std::vector<int32_t> rank;321 322    std::fill(dst_blk_pos, dst_blk_pos + 4*n_blocks*n_ns, 0);323 324    for (int64_t s = 0; s < n_ns; ++s) {325        // ubatch index s*n_tps belongs to this stream; ask which cells array it uses326        const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0];327        const auto & cells = get_mem_idx()->get_cells(seq_of_stream);328 329        int32_t * cur_cell_blk  = dst_cell_blk  + s*n_kv;330        int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks);331 332        std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0);333 334        bid_idx  .clear();335        bid_cell .clear();336        bid_slot0.clear();337 338        int n_seq_present = 0;339 340        for (int sq = 0; sq < LLAMA_MAX_SEQ && n_seq_present < 2; ++sq) {341            if (cells.seq_pos_min(sq) >= 0) {342                n_seq_present++;343            }344        }345 346        const bool one_seq = n_seq_present <= 1;347 348        // a cell no block covers needs its own -inf, which a per-block bias cannot carry349        // every cache path keeps the position below the cell window, so this stays false350        bool oor = false;351 352        bool dup = false;353 354        bool ranked = false;355 356        auto group_cells = [&]() {357            // -1 means no usable block: an incomplete or short group cannot be pooled358            std::fill(blk_of.begin(),   blk_of.end(),   -1);359            std::fill(cell_grp.begin(), cell_grp.end(), -1);360            std::fill(grp_head.begin(), grp_head.end(), -1);361 362            grp_next .clear();363            grp_first.clear();364            grp_slot0.clear();365            grp_slots.clear();366            grp_bid  .clear();367 368            oor = false;369            dup = false;370 371            for (int64_t j = 0; j < n_kv; ++j) {372                if (cells.is_empty(j)) {373                    continue;374                }375 376                const int64_t idx = ranked ? rank[j] : cells.pos_get(j);377                const int64_t pb  = idx/r;378 379                if (pb >= n_blocks) {380                    oor = true;381                    continue;382                }383 384                int32_t g = -1;385 386                for (int32_t c = grp_head[pb]; c >= 0; c = grp_next[c]) {387                    if (one_seq || cells.seq_get_all((uint32_t) grp_first[c]) == cells.seq_get_all((uint32_t) j)) {388                        g = c;389                        break;390                    }391                }392 393                if (g < 0) {394                    g = (int32_t) grp_first.size();395 396                    grp_next .push_back(grp_head[pb]);397                    grp_first.push_back((int32_t) j);398                    grp_slot0.push_back(-1);399                    grp_slots.push_back(0);400                    grp_bid  .push_back(-1);401 402                    grp_head[pb] = g;403                }404 405                const uint64_t bit = uint64_t(1) << (idx%r);406 407                dup |= (grp_slots[g] & bit) != 0;408 409                cell_grp[j]   = g;410                grp_slots[g] |= bit;411 412                if (idx%r == 0) {413                    grp_slot0[g] = (int32_t) j;414                }415            }416        };417 418        group_cells();419 420        // mrope repeats one position across an image, so rank cells instead of using the position421        if (dup && ubatch->is_pos_2d() && one_seq) {422            order.clear();423            order.reserve(n_kv);424 425            for (int64_t j = 0; j < n_kv; ++j) {426                if (!cells.is_empty(j)) {427                    order.push_back((int32_t) j);428                }429            }430 431            // same total order the mrope causal mask uses: pos, then ext.y, then ext.x432            std::sort(order.begin(), order.end(), [&cells](int32_t a, int32_t b) {433                const llama_pos pa = cells.pos_get(a);434                const llama_pos pb = cells.pos_get(b);435 436                if (pa != pb) {437                    return pa < pb;438                }439 440                const auto & ea = cells.ext_get(a);441 442                return cells.ext_get(b).is_2d_gt(ea.x, ea.y);443            });444 445            rank.assign(n_kv, -1);446 447            for (int64_t k = 0; k < (int64_t) order.size(); ++k) {448                rank[order[k]] = (int32_t) k;449            }450 451            ranked = true;452 453            group_cells();454        }455 456        GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window");457 458        int32_t n_bid = 0;459 460        for (int64_t pb = 0; pb < n_blocks; ++pb) {461            for (int32_t g = grp_head[pb]; g >= 0; g = grp_next[g]) {462                if (grp_slots[g] != slots_full) {463                    continue;464                }465 466                grp_bid[g] = n_bid++;467 468                bid_idx  .push_back((int32_t) (pb*r));469                bid_cell .push_back(grp_first[g]);470                bid_slot0.push_back(grp_slot0[g]);471            }472        }473 474        GGML_ASSERT(n_bid <= n_blocks);475 476        for (int32_t b = 0; b < n_bid; ++b) {477            int32_t sec_pos[4] = { bid_idx[b], bid_idx[b], bid_idx[b], bid_idx[b] };478 479            if (ranked) {480                const int32_t   c = bid_slot0[b];481                const llama_pos p = cells.pos_get(c);482                const auto &    e = cells.ext_get(c);483 484                sec_pos[0] = p;485                sec_pos[1] = e.y;486                sec_pos[2] = e.x;487                sec_pos[3] = p;488            }489 490            for (int64_t sec = 0; sec < 4; ++sec) {491                dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = sec_pos[sec];492            }493        }494 495        // unpooled cells all point at one spare block. a spare block exists only when some496        // cell is unpooled: n_bid == n_blocks means every cell sits in a full block.497        const bool     have_dead = n_bid < n_blocks;498        const int32_t  dead_bid  = have_dead ? n_bid : n_blocks - 1;499 500        for (int64_t j = 0; j < n_kv; ++j) {501            const int32_t g = cell_grp[j];502 503            blk_of[j] = g < 0 ? -1 : grp_bid[g];504 505            if (blk_of[j] >= 0) {506                const int64_t idx = ranked ? rank[j] : cells.pos_get(j);507 508                cur_blk_cells[blk_of[j]*r + (idx%r)] = (int32_t) j;509            }510 511            cur_cell_blk[j] = blk_of[j] < 0 ? dead_bid : blk_of[j];512        }513 514        for (int64_t ii = 0; ii < n_tps; ++ii) {515            const int64_t      i      = s*n_tps + ii;516            const llama_seq_id seq_id = ubatch->seq_id[i][0];517 518            int64_t q = ubatch->pos[i];519 520            if (ranked) {521                const llama_pos qt = ubatch->pos[i];522                const llama_pos qy = ubatch->pos[i + n_tokens];523                const llama_pos qx = ubatch->pos[i + n_tokens*2];524 525                int64_t lo = 0;526                int64_t hi = (int64_t) order.size();527 528                while (lo < hi) {529                    const int64_t   mid = (lo + hi)/2;530                    const int32_t   c   = order[mid];531                    const llama_pos pc  = cells.pos_get(c);532 533                    if (pc < qt || (pc == qt && !cells.ext_get(c).is_2d_gt(qx, qy))) {534                        lo = mid + 1;535                    } else {536                        hi = mid;537                    }538                }539 540                q = lo - 1;541            }542 543            // the tail is an incomplete block and is always visible, as in the reference544            const int64_t tail_start = (q + 1)/r*r;545 546            if (blk_bias) {547                // a block sits wholly inside or outside the tail, so one value covers it548                // the caller adds the attention mask, which drops empty, foreign and future cells549                float * cur_blk_bias = dst_bias + i*n_blocks;550 551                for (int64_t b = 0; b < n_blocks; ++b) {552                    if (b >= n_bid || !cells.seq_has((uint32_t) bid_cell[b], seq_id)) {553                        cur_blk_bias[b] = -INFINITY;554                        continue;555                    }556 557                    // finite, so it can never meet a -inf and produce a nan558                    cur_blk_bias[b] = bid_idx[b] >= tail_start ? 1e9f : 0.0f;559                }560 561                // the spare block holds the unpooled cells, which are the incomplete tail, so562                // it gets the tail value. it must stay finite: a sequence with fewer than563                // `ratio` cells owns no full block, and a row of -inf only gives a nan.564                if (have_dead) {565                    cur_blk_bias[dead_bid] = 1e9f;566                }567 568                continue;569            }570 571            float * cur_bias = dst_bias + i*n_kv;572 573            for (int64_t j = 0; j < n_kv; ++j) {574                float v = -INFINITY;575 576                if (!cells.is_empty(j) && cells.seq_has(j, seq_id)) {577                    const int64_t idx = ranked ? rank[j] : cells.pos_get(j);578 579                    if (idx <= q) {580                        // finite, so it can never meet a -inf and produce a nan581                        v = idx >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f);582                    }583                }584 585                cur_bias[j] = v;586            }587        }588    }589}590 591//592// llama_memory_hybrid_idx_context593//594 595// streams in each ubatch's slot info, matching get_k/get_v's `ns`596static std::vector<uint32_t> llama_memory_hybrid_idx_ns(const llama_kv_cache::slot_info_vec_t & sinfos) {597    std::vector<uint32_t> res;598    res.reserve(sinfos.size());599 600    for (const auto & sinfo : sinfos) {601        res.push_back(sinfo.s1 - sinfo.s0 + 1);602    }603 604    return res;605}606 607llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) :608    llama_memory_hybrid_context(status) {}609 610llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem) :611    llama_memory_hybrid_context(mem),612    mem(mem),613    // graph reservation walks a full context, and qwen4exp builds the sparse attention only when this is set614    // without it the reserved worst case is the dense graph, so ggml-alloc must grow the buffer on the first decode615    ns_ubatch(mem->get_mem_idx() == nullptr ?616        std::vector<uint32_t>() : std::vector<uint32_t>{ mem->get_mem_idx()->get_n_stream() }),617    ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :618        new llama_kv_cache_context(mem->get_mem_idx())) {}619 620llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(621        llama_memory_hybrid_idx * mem,622                  llama_context * lctx,623                           bool   optimize) :624    llama_memory_hybrid_context(mem, lctx, optimize),625    mem(mem),626    // update() applies a pending cross-stream seq_cp, else the copy keeps stale indexer keys627    ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :628        mem->get_mem_idx()->init_update(lctx, optimize)) {}629 630llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(631        llama_memory_hybrid_idx * mem,632                slot_info_vec_t   sinfos_attn,633                slot_info_vec_t   sinfos_idx,634      std::vector<llama_ubatch>   ubatches) :635    // note: the base copies the ubatches; ctx_idx gets a copy of its own636    llama_memory_hybrid_context(mem, std::move(sinfos_attn), ubatches),637    mem(mem),638    ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)),639    ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :640        new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {}641 642bool llama_memory_hybrid_idx_context::next() {643    if (ctx_idx) {644        ctx_idx->next();645    }646 647    ++i_cur;648 649    return llama_memory_hybrid_context::next();650}651 652bool llama_memory_hybrid_idx_context::apply() {653    bool res = llama_memory_hybrid_context::apply();654 655    if (ctx_idx) {656        res = res & ctx_idx->apply();657    }658 659    return res;660}661 662const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const {663    return static_cast<const llama_kv_cache_context *>(ctx_idx.get());664}665 666uint32_t llama_memory_hybrid_idx_context::get_n_stream() const {667    GGML_ASSERT(i_cur < ns_ubatch.size());668 669    return ns_ubatch[i_cur];670}671 672void llama_memory_hybrid_idx_context::set_input_qsa(673        ggml_tensor * cell_blk,674        ggml_tensor * blk_cells,675        ggml_tensor * blk_pos,676        ggml_tensor * bias,677        const llama_ubatch * ubatch,678        uint32_t ratio,679        bool blk_bias) const {680    GGML_ASSERT(mem != nullptr);681 682    mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias);683}684