diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e..55054bb9d2 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -29,8 +29,10 @@ llama_memory_hybrid::llama_memory_hybrid( bool unified, /* layer filters */ const layer_filter_cb & filter_attn, - const layer_filter_cb & filter_recr) : + const layer_filter_cb & filter_recr, + const layer_filter_cb & filter_idx) : hparams(model.hparams), + hparams_idx(model.hparams), mem_attn(new llama_kv_cache( model, model.hparams, @@ -62,7 +64,20 @@ llama_memory_hybrid::llama_memory_hybrid( filter_recr == nullptr ? [&](int32_t il) { return hparams.is_recr(il); } : filter_recr - )) {} + )), + mem_idx(filter_idx == nullptr ? nullptr : [&] { + // MQA with a single key head of indexer_head_size, the same shaping + // llama_kv_cache_dsa applies to its lightning-indexer cache + std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + + LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); + + return new llama_kv_cache( + model, hparams_idx, type_k, type_v, v_trans, offload, unified, + kv_size, n_seq_max, n_pad, n_swa, swa_type, + nullptr, filter_idx, nullptr, nullptr); + }()) {} llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { @@ -109,6 +124,14 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba } // prepare the attention cache + llama_kv_cache::slot_info_vec_t heads_idx; + if (mem_idx) { + heads_idx = mem_idx->prepare(ubatches); + if (heads_idx.empty()) { + break; + } + } + auto heads_attn = mem_attn->prepare(ubatches); if (heads_attn.empty()) { LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__); @@ -116,7 +139,7 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba } return std::make_unique( - this, std::move(heads_attn), std::move(ubatches)); + this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches)); } while(false); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); @@ -137,6 +160,7 @@ bool llama_memory_hybrid::get_can_shift() const { void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); + if (mem_idx) mem_idx->clear(data); mem_recr->clear(data); } @@ -146,26 +170,31 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 if (!mem_recr->seq_rm(seq_id, p0, p1)) { return false; } + if (mem_idx) mem_idx->seq_rm(seq_id, p0, p1); return mem_attn->seq_rm(seq_id, p0, p1); } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); + if (mem_idx) mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); } void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { mem_attn->seq_keep(seq_id); + if (mem_idx) mem_idx->seq_keep(seq_id); mem_recr->seq_keep(seq_id); } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { mem_attn->seq_add(seq_id, p0, p1, shift); + if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { mem_attn->seq_div(seq_id, p0, p1, d); + if (mem_idx) mem_idx->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); } @@ -205,6 +234,10 @@ llama_kv_cache * llama_memory_hybrid::get_mem_attn() const { return mem_attn.get(); } +llama_kv_cache * llama_memory_hybrid::get_mem_idx() const { + return mem_idx.get(); +} + llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const { return mem_recr.get(); } @@ -229,11 +262,14 @@ llama_memory_hybrid_context::llama_memory_hybrid_context( llama_memory_hybrid_context::llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, std::vector ubatches) : ubatches(std::move(ubatches)), // note: here we copy the ubatches. not sure if this is ideal ctx_attn(new llama_kv_cache_context(mem->get_mem_attn(), std::move(sinfos_attn), this->ubatches)), ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), this->ubatches)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -242,6 +278,7 @@ bool llama_memory_hybrid_context::next() { ctx_attn->next(); ctx_recr->next(); + if (ctx_idx) ctx_idx->next(); if (++i_next >= ubatches.size()) { return false; @@ -257,6 +294,7 @@ bool llama_memory_hybrid_context::apply() { res = res & ctx_attn->apply(); res = res & ctx_recr->apply(); + if (ctx_idx) res = res & ctx_idx->apply(); return res; } @@ -277,3 +315,7 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const { const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const llama_kv_cache_context * llama_memory_hybrid_context::get_idx() const { + return static_cast(ctx_idx.get()); +} diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb749..2c03ee7ca5 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -39,7 +39,12 @@ public: bool unified, /* layer filters */ const layer_filter_cb & filter_attn = nullptr, - const layer_filter_cb & filter_recr = nullptr); + const layer_filter_cb & filter_recr = nullptr, + /* optional per-token indexer key cache, for hybrid + models whose attention layers are sparse. absent + unless filter_idx is given, so every existing + architecture is unaffected. */ + const layer_filter_cb & filter_idx = nullptr); ~llama_memory_hybrid() = default; @@ -82,12 +87,18 @@ public: llama_kv_cache * get_mem_attn() const; llama_memory_recurrent * get_mem_recr() const; + llama_kv_cache * get_mem_idx() const; // nullptr when the model has no indexer private: const llama_hparams & hparams; + // geometry for the indexer cache: MQA with a single key head of + // indexer_head_size, mirroring how llama_kv_cache_dsa builds its own + llama_hparams hparams_idx; + const std::unique_ptr mem_attn; const std::unique_ptr mem_recr; + const std::unique_ptr mem_idx; }; class llama_memory_hybrid_context : public llama_memory_context_i { @@ -110,6 +121,7 @@ public: llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, std::vector ubatches); ~llama_memory_hybrid_context() = default; @@ -126,6 +138,7 @@ public: const llama_kv_cache_context * get_attn() const; const llama_memory_recurrent_context * get_recr() const; + const llama_kv_cache_context * get_idx() const; // nullptr without an indexer private: // the index of the next ubatch to process @@ -135,6 +148,7 @@ private: const llama_memory_context_ptr ctx_attn; const llama_memory_context_ptr ctx_recr; + const llama_memory_context_ptr ctx_idx; // null unless the model has an indexer const llama_memory_status status; };