ikawrakow/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-02-25 07:34:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Avoid allocating MHA KV cache when MLA is turned on

Browse Source
This commit is contained in:
Saood Karim
2025-02-08 12:47:52 -06:00
parent 3aaf602da5
commit f0227c4bfa
1 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
src/llama.cpp
View File

@@ -3173,8 +3173,17 @@ static bool llama_kv_cache_init(
const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(i) + hparams.n_embd_v_s();
struct ggml_context * ctx = offload ? ctx_map.at(model.buft_layer[i].buft) : cache.ctxs.front();
ggml_tensor * k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size);
ggml_tensor * v = ggml_new_tensor_1d(ctx, type_v, n_embd_v_gqa*kv_size);
ggml_tensor * k;
ggml_tensor * v;
if (cparams.mla_attn) {
k = ggml_new_tensor_1d(ctx, type_k, 1);
v = ggml_new_tensor_1d(ctx, type_v, 1);
}
else {
k = ggml_new_tensor_1d(ctx, type_k, n_embd_k_gqa*kv_size);
v = ggml_new_tensor_1d(ctx, type_v, n_embd_v_gqa*kv_size);
}
ggml_format_name(k, "cache_k_l%d", i);
ggml_format_name(v, "cache_v_l%d", i);
cache.k_l.push_back(k);