Skip to content

Commit b6c4e55

Browse files
ggerganovjordankanter
ggerganov
authored and
jordankanter
committed
llama : refactor k-shift implementation + KV defragmentation (ggml-org#5691)
* llama : refactor k-shift implementation ggml-ci * llama : rename llama_kv_cache_seq_shift to llama_kv_cache_seq_add * llama : cont k-shift refactoring + normalize type names ggml-ci * minor : fix MPI builds * llama : reuse n_rot from the build context ggml-ci * llama : revert enum name changes from this PR ggml-ci * llama : update llama_rope_type * llama : add comment about rope values * llama : fix build * passkey : apply kv cache updates explicitly ggml-ci * llama : change name to llama_kv_cache_update() * llama : add llama_kv_cache_seq_pos_max() * passkey : fix llama_kv_cache_seq_pos_max() usage * llama : some llama_kv_cell simplifications * llama : add llama_kv_cache_compress (EXPERIMENTAL) * llama : add alternative KV cache merging (EXPERIMENTAL) * llama : add llama_kv_cache_defrag * llama : comments * llama : remove llama_kv_cache_compress will add in a separate PR ggml-ci * llama : defragment via non-overlapping moves * llama : ggml_graph based defrag implementation ggml-ci * llama : switch the loop order in build_defrag * llama : add comments
1 parent 4e52b58 commit b6c4e55

File tree

6 files changed

+646
-304
lines changed

6 files changed

+646
-304
lines changed

‎examples/infill/infill.cpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -447,8 +447,8 @@ int main(int argc, char ** argv) {
447447
LOG("context full, swapping: n_past = %d, n_left = %d, n_ctx = %d, n_keep = %d, n_discard = %d\n",
448448
n_past, n_left, n_ctx, params.n_keep, n_discard);
449449

450-
llama_kv_cache_seq_rm (ctx, 0, params.n_keep + 1 , params.n_keep + n_discard + 1);
451-
llama_kv_cache_seq_shift(ctx, 0, params.n_keep + 1 + n_discard, n_past, -n_discard);
450+
llama_kv_cache_seq_rm (ctx, 0, params.n_keep + 1 , params.n_keep + n_discard + 1);
451+
llama_kv_cache_seq_add(ctx, 0, params.n_keep + 1 + n_discard, n_past, -n_discard);
452452

453453
n_past -= n_discard;
454454

‎examples/main/main.cpp

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -548,8 +548,8 @@ int main(int argc, char ** argv) {
548548
LOG("context full, swapping: n_past = %d, n_left = %d, n_ctx = %d, n_keep = %d, n_discard = %d\n",
549549
n_past, n_left, n_ctx, params.n_keep, n_discard);
550550

551-
llama_kv_cache_seq_rm (ctx, 0, params.n_keep , params.n_keep + n_discard);
552-
llama_kv_cache_seq_shift(ctx, 0, params.n_keep + n_discard, n_past, -n_discard);
551+
llama_kv_cache_seq_rm (ctx, 0, params.n_keep , params.n_keep + n_discard);
552+
llama_kv_cache_seq_add(ctx, 0, params.n_keep + n_discard, n_past, -n_discard);
553553

554554
n_past -= n_discard;
555555

@@ -576,9 +576,9 @@ int main(int argc, char ** argv) {
576576
LOG("div: [%6d, %6d] / %6d -> [%6d, %6d]\n", ga_i + ib*bd, ga_i + ib*bd + ga_w, ga_n, (ga_i + ib*bd)/ga_n, (ga_i + ib*bd + ga_w)/ga_n);
577577
LOG("shift: [%6d, %6d] + %6d -> [%6d, %6d]\n", ga_i + ib*bd + ga_w, n_past + ib*bd, dd, ga_i + ib*bd + ga_w + dd, n_past + ib*bd + dd);
578578

579-
llama_kv_cache_seq_shift(ctx, 0, ga_i, n_past, ib*bd);
580-
llama_kv_cache_seq_div (ctx, 0, ga_i + ib*bd, ga_i + ib*bd + ga_w, ga_n);
581-
llama_kv_cache_seq_shift(ctx, 0, ga_i + ib*bd + ga_w, n_past + ib*bd, dd);
579+
llama_kv_cache_seq_add(ctx, 0, ga_i, n_past, ib*bd);
580+
llama_kv_cache_seq_div(ctx, 0, ga_i + ib*bd, ga_i + ib*bd + ga_w, ga_n);
581+
llama_kv_cache_seq_add(ctx, 0, ga_i + ib*bd + ga_w, n_past + ib*bd, dd);
582582

583583
n_past -= bd;
584584

‎examples/passkey/passkey.cpp

Lines changed: 15 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ int main(int argc, char ** argv) {
126126
const int n_batch = ctx_params.n_batch;
127127
const int n_batch_grp = ctx_params.n_batch/n_grp;
128128

129-
LOG_TEE("\n%s: n_len = %d, n_ctx = %d, n_kv_req = %d, n_grp = %d, n_batch = %d\n", __func__, n_len, n_ctx, n_kv_req, n_grp, n_batch);
129+
LOG_TEE("\n%s: n_len = %d, n_ctx = %d, n_kv_req = %d, n_grp = %d, n_batch = %d, n_junk = %d, i_pos = %d\n", __func__, n_len, n_ctx, n_kv_req, n_grp, n_batch, n_junk, i_pos);
130130

131131
// print the prompt token-by-token
132132

@@ -146,10 +146,11 @@ int main(int argc, char ** argv) {
146146
const int ib = i/n_batch - 1;
147147
const int bd = n_batch_grp*(n_grp - 1);
148148

149-
llama_kv_cache_seq_shift(ctx, 0, n_past - n_batch, n_past, ib*bd);
150-
llama_kv_cache_seq_div (ctx, 0, n_past - n_batch + ib*bd, n_past + ib*bd, n_grp);
149+
llama_kv_cache_seq_add (ctx, 0, n_past - n_batch, n_past, ib*bd);
150+
llama_kv_cache_seq_div (ctx, 0, n_past - n_batch + ib*bd, n_past + ib*bd, n_grp);
151+
llama_kv_cache_update (ctx);
151152

152-
n_past -= bd;
153+
n_past = llama_kv_cache_seq_pos_max(ctx, 0) + 1;
153154
}
154155

155156
llama_batch_clear(batch);
@@ -179,10 +180,12 @@ int main(int argc, char ** argv) {
179180

180181
LOG_TEE("%s: shifting KV cache with %d\n", __func__, n_discard);
181182

182-
llama_kv_cache_seq_rm (ctx, 0, n_keep , n_keep + n_discard);
183-
llama_kv_cache_seq_shift(ctx, 0, n_keep + n_discard, n_ctx, -n_discard);
183+
llama_kv_cache_seq_rm (ctx, 0, n_keep , n_keep + n_discard);
184+
llama_kv_cache_seq_add(ctx, 0, n_keep + n_discard, n_ctx, -n_discard);
185+
llama_kv_cache_defrag (ctx);
186+
llama_kv_cache_update (ctx);
184187

185-
n_past -= n_discard;
188+
n_past = llama_kv_cache_seq_pos_max(ctx, 0) + 1;
186189

187190
llama_batch_clear(batch);
188191

@@ -208,10 +211,12 @@ int main(int argc, char ** argv) {
208211
if (n_discard > 0) {
209212
LOG_TEE("%s: shifting KV cache with %d to free space for the answer\n", __func__, n_discard);
210213

211-
llama_kv_cache_seq_rm (ctx, 0, n_keep , n_keep + n_discard);
212-
llama_kv_cache_seq_shift(ctx, 0, n_keep + n_discard, n_ctx, -n_discard);
214+
llama_kv_cache_seq_rm (ctx, 0, n_keep , n_keep + n_discard);
215+
llama_kv_cache_seq_add(ctx, 0, n_keep + n_discard, n_ctx, -n_discard);
216+
llama_kv_cache_defrag (ctx);
217+
llama_kv_cache_update (ctx);
213218

214-
n_past -= n_discard;
219+
n_past = llama_kv_cache_seq_pos_max(ctx, 0) + 1;
215220
}
216221
}
217222

‎examples/server/server.cpp

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1636,8 +1636,8 @@ struct llama_server_context
16361636
{"n_system_tokens", system_tokens.size()},
16371637
{"n_cache_tokens", slot.cache_tokens.size()}
16381638
});
1639-
llama_kv_cache_seq_rm (ctx, slot.id, n_keep , n_keep + n_discard);
1640-
llama_kv_cache_seq_shift(ctx, slot.id, n_keep + n_discard, system_tokens.size() + slot.n_past, -n_discard);
1639+
llama_kv_cache_seq_rm (ctx, slot.id, n_keep , n_keep + n_discard);
1640+
llama_kv_cache_seq_add(ctx, slot.id, n_keep + n_discard, system_tokens.size() + slot.n_past, -n_discard);
16411641

16421642
for (size_t i = n_keep + n_discard; i < slot.cache_tokens.size(); i++)
16431643
{
@@ -1941,9 +1941,9 @@ struct llama_server_context
19411941
LOG_TEE("div: [%6d, %6d] / %6d -> [%6d, %6d]\n", slot.ga_i + ib * bd, slot.ga_i + ib * bd + slot.ga_w, slot.ga_n, (slot.ga_i + ib * bd) / slot.ga_n, (slot.ga_i + ib * bd + slot.ga_w) / slot.ga_n);
19421942
LOG_TEE("shift: [%6d, %6d] + %6d -> [%6d, %6d]\n", slot.ga_i + ib * bd + slot.ga_w, slot.n_past_se + ib * bd, dd, slot.ga_i + ib * bd + slot.ga_w + dd, slot.n_past_se + ib * bd + dd);
19431943

1944-
llama_kv_cache_seq_shift(ctx, slot.id, slot.ga_i, slot.n_past_se, ib * bd);
1944+
llama_kv_cache_seq_add(ctx, slot.id, slot.ga_i, slot.n_past_se, ib * bd);
19451945
llama_kv_cache_seq_div(ctx, slot.id, slot.ga_i + ib * bd, slot.ga_i + ib * bd + slot.ga_w,slot.ga_n);
1946-
llama_kv_cache_seq_shift(ctx, slot.id, slot.ga_i + ib * bd + slot.ga_w,slot.n_past_se + ib * bd, dd);
1946+
llama_kv_cache_seq_add(ctx, slot.id, slot.ga_i + ib * bd + slot.ga_w,slot.n_past_se + ib * bd, dd);
19471947

19481948
slot.n_past_se -= bd;
19491949

0 commit comments

Comments
 (0)