Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 17:24:57 +02:00
Code Issues Packages Projects Releases Wiki Activity

llama: optional indexer key cache in llama_memory_hybrid

Browse Source 
Groundwork for qwen4exp's QSA sparse attention. Its indexer needs a per-token
key history for the full-attention layers, but a hybrid model cannot use
llama_kv_cache_dsa: that class derives from llama_memory_i rather than
llama_kv_cache, and llama_memory_hybrid constructs its attention cache
directly. No existing architecture pairs recurrent state with a sparse
indexer, so there was nothing to reuse wholesale.

llama_memory_hybrid therefore gains a third, optional cache, shaped the same
way llama_kv_cache_dsa shapes its lightning-indexer cache: a copy of hparams
with n_head_kv forced to 1 and n_embd_head_k_full set to indexer_head_size.
It is built only when a filter_idx callback is passed, which defaults to
nullptr, so every existing architecture gets exactly what it got before. The
per-sequence operations and the batch preparation forward to it under a null
check, matching how the DSA cache prepares its two caches over the same
ubatches.

test-llama-archs passes all 124 architectures at 0.00e+00, including the 12 in
the hybrid family that share this code. The qwen4exp fixtures are unchanged:
same logits against vLLM, and chunked evaluation still bit-identical to
single-shot.
This commit is contained in:
Daniel Han
2026-08-25 18:16:16 +00:00
parent ddf0980e52
commit ff24f3874f
2 changed files with 60 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/llama-memory-hybrid.cpp
+45 -3
View File
@@ -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<llama_memory_hybrid_context>(
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_hybrid_context>(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<llama_ubatch> 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<const llama_memory_recurrent_context *>(ctx_recr.get());
}
const llama_kv_cache_context * llama_memory_hybrid_context::get_idx() const {
return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
}
src/llama-memory-hybrid.h
+15 -1
View File
@@ -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<llama_kv_cache> mem_attn;
const std::unique_ptr<llama_memory_recurrent> mem_recr;
const std::unique_ptr<llama_kv_cache> 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<llama_ubatch> 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;
};