llama : allocate indexer cache only in "full" indexer layers (#26474)

Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
This commit is contained in:
fairydreaming
2026-08-03 14:56:30 +02:00
committed by GitHub
co-authored by Stanisław Szymczyk
parent 96278e39fc
commit 563dec81c1
3 changed files with 11 additions and 7 deletions
+4 -3
View File
@@ -23,7 +23,8 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_filter_cb & filter_mla,
const layer_filter_cb & filter_lid,
const layer_reuse_cb & reuse) :
hparams_lid(model.hparams), n_stream(unified ? 1 : n_seq_max) {
@@ -32,7 +33,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
kv_mla = std::make_unique<llama_kv_cache>(
model, model.hparams, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
n_swa, swa_type, nullptr, filter, reuse, nullptr);
n_swa, swa_type, nullptr, filter_mla, reuse, nullptr);
// we use llama_kv_cache for caching indexer keys
// by hand-tweaking some hparams we fool it to create
@@ -49,7 +50,7 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
kv_lid = std::make_unique<llama_kv_cache>(
model, hparams_lid, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max, n_pad,
n_swa, swa_type, nullptr, filter, reuse, nullptr);
n_swa, swa_type, nullptr, filter_lid, reuse, nullptr);
}
void llama_kv_cache_dsa::clear(bool data) {
+2 -1
View File
@@ -26,7 +26,8 @@ public:
uint32_t n_pad,
uint32_t n_swa,
llama_swa_type swa_type,
const layer_filter_cb & filter,
const layer_filter_cb & filter_mla,
const layer_filter_cb & filter_lid,
const layer_reuse_cb & reuse);
~llama_kv_cache_dsa() = default;
+5 -3
View File
@@ -2101,10 +2101,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
} else {
// Main context: DSA cache for the trunk layers only - the nextn
// layer(s) are never attended by the trunk graph.
llama_kv_cache::layer_filter_cb filter = nullptr;
llama_kv_cache::layer_filter_cb filter_mla = nullptr;
if (hparams.n_layer_nextn > 0) {
filter = [&](uint32_t il) { return il < hparams.n_layer(); };
filter_mla = [&](uint32_t il) { return il < hparams.n_layer(); };
}
llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return il < hparams.n_layer() && (arch != LLM_ARCH_GLM_DSA || hparams.is_indexer_full(il)); };
res = new llama_kv_cache_dsa(
*this,
@@ -2118,7 +2119,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
1,
hparams.n_swa,
hparams.swa_type,
filter,
filter_mla,
filter_lid,
nullptr);
}
} break;