llama : remove implicit recurrent state rollbacks

compilade · compilade · commit 8006f3b3c83d · 2024-11-24T20:35:30.000-05:00
diff --git a/common/common.cpp b/common/common.cpp
@@ -966,7 +966,7 @@ struct common_init_result common_init_from_params(common_params & params) {
         if (llama_model_has_decoder(model)) {
             llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch), 0, 0));
         }
-        llama_past_clear(lctx);
+        llama_kv_cache_clear(lctx);
         llama_synchronize(lctx);
         llama_perf_context_reset(lctx);
     }
diff --git a/examples/batched-bench/batched-bench.cpp b/examples/batched-bench/batched-bench.cpp
@@ -133,7 +133,7 @@ int main(int argc, char ** argv) {
 
                 const auto t_pp_start = ggml_time_us();
 
-                llama_past_clear(ctx);
+                llama_kv_cache_clear(ctx);
 
                 if (!decode_helper(ctx, batch, ctx_params.n_batch)) {
                     LOG_ERR("%s: llama_decode() failed\n", __func__);
@@ -142,7 +142,7 @@ int main(int argc, char ** argv) {
 
                 if (is_pp_shared) {
                     for (int32_t i = 1; i < pl; ++i) {
-                        llama_past_seq_cp(ctx, 0, i, -1, -1);
+                        llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);
                     }
                 }
 
diff --git a/examples/batched.swift/Sources/main.swift b/examples/batched.swift/Sources/main.swift
@@ -111,7 +111,7 @@ if llama_decode(context, batch) != 0 {
 }
 
 for i in 1 ..< n_parallel {
-    llama_past_seq_cp(context, 0, Int32(i), -1, -1)
+    llama_kv_cache_seq_cp(context, 0, Int32(i), -1, -1)
 }
 
 if n_parallel > 1 {
diff --git a/examples/batched/batched.cpp b/examples/batched/batched.cpp
@@ -138,7 +138,7 @@ int main(int argc, char ** argv) {
     //// assign the system KV cache to all parallel sequences
     //// this way, the parallel sequences will "reuse" the prompt tokens without having to copy them
     //for (int32_t i = 1; i < n_parallel; ++i) {
-    //    llama_past_seq_cp(ctx, 0, i, -1, -1);
+    //    llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);
     //}
 
     if (n_parallel > 1) {
diff --git a/examples/cvector-generator/cvector-generator.cpp b/examples/cvector-generator/cvector-generator.cpp
@@ -338,7 +338,7 @@ static bool cb_eval(struct ggml_tensor * t, bool ask, void * user_data) {
 }
 
 static bool get_hidden_layers(llama_context * ctx, std::vector<llama_token> & tokens) {
-    llama_past_clear(ctx);
+    llama_kv_cache_clear(ctx);
     if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size(), 0, 0))) {
         fprintf(stderr, "%s : failed to eval\n", __func__);
         return false;
diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp
@@ -37,7 +37,7 @@ static void batch_decode(llama_context * ctx, llama_batch & batch, float * outpu
     const struct llama_model * model = llama_get_model(ctx);
 
     // clear previous kv_cache values (irrelevant for embeddings)
-    llama_past_clear(ctx);
+    llama_kv_cache_clear(ctx);
 
     // run model
     LOG_INF("%s: n_tokens = %d, n_seq = %d\n", __func__, batch.n_tokens, n_seq);
diff --git a/examples/gritlm/gritlm.cpp b/examples/gritlm/gritlm.cpp
@@ -44,7 +44,7 @@ static std::vector<std::vector<float>> encode(llama_context * ctx, const std::ve
         }
 
         // clear previous kv_cache values (irrelevant for embeddings)
-        llama_past_clear(ctx);
+        llama_kv_cache_clear(ctx);
         llama_set_embeddings(ctx, true);
         llama_set_causal_attn(ctx, false);
 
@@ -99,7 +99,7 @@ static std::string generate(llama_context * ctx, llama_sampler * smpl, const std
     const llama_model * model = llama_get_model(ctx);
     llama_token eos_token = llama_token_eos(model);
 
-    llama_past_clear(ctx);
+    llama_kv_cache_clear(ctx);
     llama_set_embeddings(ctx, false);
     llama_set_causal_attn(ctx, true);
 
diff --git a/examples/imatrix/imatrix.cpp b/examples/imatrix/imatrix.cpp
@@ -494,7 +494,7 @@ static bool compute_imatrix(llama_context * ctx, const common_params & params) {
         const auto t_start = std::chrono::high_resolution_clock::now();
 
         // clear the KV cache
-        llama_past_clear(ctx);
+        llama_kv_cache_clear(ctx);
 
         for (int j = 0; j < num_batches; ++j) {
             const int batch_start = start + j * n_batch;
diff --git a/examples/infill/infill.cpp b/examples/infill/infill.cpp
@@ -375,8 +375,8 @@ int main(int argc, char ** argv) {
                 LOG_DBG("context full, swapping: n_past = %d, n_left = %d, n_ctx = %d, n_keep = %d, n_discard = %d\n",
                     n_past, n_left, n_ctx, params.n_keep, n_discard);
 
-                llama_past_seq_rm (ctx, 0, params.n_keep + 1            , params.n_keep + n_discard + 1);
-                llama_past_seq_add(ctx, 0, params.n_keep + 1 + n_discard, n_past, -n_discard);
+                llama_kv_cache_seq_rm (ctx, 0, params.n_keep + 1            , params.n_keep + n_discard + 1);
+                llama_kv_cache_seq_add(ctx, 0, params.n_keep + 1 + n_discard, n_past, -n_discard);
 
                 n_past -= n_discard;
 
diff --git a/examples/llama-bench/llama-bench.cpp b/examples/llama-bench/llama-bench.cpp
@@ -1566,7 +1566,7 @@ int main(int argc, char ** argv) {
 
         test t(inst, lmodel, ctx);
 
-        llama_past_clear(ctx);
+        llama_kv_cache_clear(ctx);
 
         // cool off before the test
         if (params.delay) {
@@ -1606,7 +1606,7 @@ int main(int argc, char ** argv) {
         }
 
         for (int i = 0; i < params.reps; i++) {
-            llama_past_clear(ctx);
+            llama_kv_cache_clear(ctx);
 
             uint64_t t_start = get_time_ns();
 
diff --git a/examples/llama.android/llama/src/main/cpp/llama-android.cpp b/examples/llama.android/llama/src/main/cpp/llama-android.cpp
@@ -194,7 +194,7 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model(
         }
 
         batch->logits[batch->n_tokens - 1] = true;
-        llama_past_clear(context);
+        llama_kv_cache_clear(context);
 
         const auto t_pp_start = ggml_time_us();
         if (llama_decode(context, *batch) != 0) {
@@ -206,7 +206,7 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model(
 
         LOGi("Benchmark text generation (tg)");
 
-        llama_past_clear(context);
+        llama_kv_cache_clear(context);
         const auto t_tg_start = ggml_time_us();
         for (i = 0; i < tg; i++) {
 
@@ -223,7 +223,7 @@ Java_android_llama_cpp_LLamaAndroid_bench_1model(
 
         const auto t_tg_end = ggml_time_us();
 
-        llama_past_clear(context);
+        llama_kv_cache_clear(context);
 
         const auto t_pp = double(t_pp_end - t_pp_start) / 1000000.0;
         const auto t_tg = double(t_tg_end - t_tg_start) / 1000000.0;
@@ -446,5 +446,5 @@ Java_android_llama_cpp_LLamaAndroid_completion_1loop(
 extern "C"
 JNIEXPORT void JNICALL
 Java_android_llama_cpp_LLamaAndroid_kv_1cache_1clear(JNIEnv *, jobject, jlong context) {
-    llama_past_clear(reinterpret_cast<llama_context *>(context));
+    llama_kv_cache_clear(reinterpret_cast<llama_context *>(context));
 }
diff --git a/examples/llama.swiftui/llama.cpp.swift/LibLlama.swift b/examples/llama.swiftui/llama.cpp.swift/LibLlama.swift
@@ -209,7 +209,7 @@ actor LlamaContext {
             }
             batch.logits[Int(batch.n_tokens) - 1] = 1 // true
 
-            llama_past_clear(context)
+            llama_kv_cache_clear(context)
 
             let t_pp_start = ggml_time_us()
 
@@ -222,7 +222,7 @@ actor LlamaContext {
 
             // bench text generation
 
-            llama_past_clear(context)
+            llama_kv_cache_clear(context)
 
             let t_tg_start = ggml_time_us()
 
@@ -241,7 +241,7 @@ actor LlamaContext {
 
             let t_tg_end = ggml_time_us()
 
-            llama_past_clear(context)
+            llama_kv_cache_clear(context)
 
             let t_pp = Double(t_pp_end - t_pp_start) / 1000000.0
             let t_tg = Double(t_tg_end - t_tg_start) / 1000000.0
@@ -291,7 +291,7 @@ actor LlamaContext {
     func clear() {
         tokens_list.removeAll()
         temporary_invalid_cchars.removeAll()
-        llama_past_clear(context)
+        llama_kv_cache_clear(context)
     }
 
     private func tokenize(text: String, add_bos: Bool) -> [llama_token] {
diff --git a/examples/lookahead/lookahead.cpp b/examples/lookahead/lookahead.cpp
@@ -93,7 +93,7 @@ int main(int argc, char ** argv) {
     llama_decode(ctx, llama_batch_get_one(&inp.back(),           1, n_input - 1, 0));
 
     for (int s = 1; s < W + G + 1; ++s) {
-        llama_past_seq_cp(ctx, 0, s, -1, -1);
+        llama_kv_cache_seq_cp(ctx, 0, s, -1, -1);
     }
 
     const auto t_enc_end = ggml_time_us();
@@ -436,17 +436,17 @@ int main(int argc, char ** argv) {
         // KV cache management
         // if no verification token matched, we simply remove all cells from this batch -> no fragmentation
         // FIXME: recurrent and hybrid models
-        llama_past_seq_rm(ctx, -1, n_past, -1);
+        llama_kv_cache_seq_rm(ctx, -1, n_past, -1);
 
         if (seq_id_best != 0) {
             // if a verification token matched, we keep the best sequence and remove the rest
             // this leads to some KV cache fragmentation
-            llama_past_seq_keep(ctx, seq_id_best);
-            llama_past_seq_cp  (ctx, seq_id_best, 0, -1, -1);
-            llama_past_seq_rm  (ctx, seq_id_best,    -1, -1);
+            llama_kv_cache_seq_keep(ctx, seq_id_best);
+            llama_kv_cache_seq_cp  (ctx, seq_id_best, 0, -1, -1);
+            llama_kv_cache_seq_rm  (ctx, seq_id_best,    -1, -1);
 
             for (int s = 1; s < W + G + 1; ++s) {
-                llama_past_seq_cp(ctx, 0, s, -1, -1);
+                llama_kv_cache_seq_cp(ctx, 0, s, -1, -1);
             }
         }
     }
diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp
@@ -191,7 +191,7 @@ int main(int argc, char ** argv){
         // KV cache management
         // clean the cache of draft tokens that weren't accepted
         // FIXME: recurrent and hybrid models
-        llama_past_seq_rm(ctx, 0, n_past, -1);
+        llama_kv_cache_seq_rm(ctx, 0, n_past, -1);
 
         common_batch_clear(batch_tgt);
         common_batch_add(batch_tgt, draft[0], n_past, { 0 }, true);
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
@@ -332,10 +332,6 @@ int main(int argc, char ** argv) {
             }
             n_matching_session_tokens++;
         }
-
-        // remove any "future" tokens that we might have inherited from the previous session
-        n_matching_session_tokens = llama_past_seq_rm(ctx, -1, n_matching_session_tokens, -1);
-
         if (params.prompt.empty() && n_matching_session_tokens == embd_inp.size()) {
             LOG_INF("%s: using full prompt from session file\n", __func__);
         } else if (n_matching_session_tokens >= embd_inp.size()) {
@@ -347,6 +343,9 @@ int main(int argc, char ** argv) {
             LOG_INF("%s: session file matches %zu / %zu tokens of prompt\n",
                     __func__, n_matching_session_tokens, embd_inp.size());
         }
+
+        // remove any "future" tokens that we might have inherited from the previous session
+        llama_kv_cache_seq_rm(ctx, -1, n_matching_session_tokens, -1);
     }
 
     LOG_DBG("recalculate the cached logits (check): embd_inp.size() %zu, n_matching_session_tokens %zu, embd_inp.size() %zu, session_tokens.size() %zu\n",
@@ -358,8 +357,6 @@ int main(int argc, char ** argv) {
         LOG_DBG("recalculate the cached logits (do): session_tokens.resize( %zu )\n", embd_inp.size() - 1);
 
         session_tokens.resize(embd_inp.size() - 1);
-    } else {
-        session_tokens.resize(n_matching_session_tokens);
     }
 
     // number of tokens to keep when resetting context
@@ -609,9 +606,9 @@ int main(int argc, char ** argv) {
                     LOG_DBG("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);
                     LOG_DBG("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);
 
-                    llama_past_seq_add(ctx, 0, ga_i,                n_past,              ib*bd);
-                    llama_past_seq_div(ctx, 0, ga_i + ib*bd,        ga_i + ib*bd + ga_w, ga_n);
-                    llama_past_seq_add(ctx, 0, ga_i + ib*bd + ga_w, n_past + ib*bd,      dd);
+                    llama_kv_cache_seq_add(ctx, 0, ga_i,                n_past,              ib*bd);
+                    llama_kv_cache_seq_div(ctx, 0, ga_i + ib*bd,        ga_i + ib*bd + ga_w, ga_n);
+                    llama_kv_cache_seq_add(ctx, 0, ga_i + ib*bd + ga_w, n_past + ib*bd,      dd);
 
                     n_past -= bd;
 
@@ -625,8 +622,6 @@ int main(int argc, char ** argv) {
             if (n_session_consumed < (int) session_tokens.size()) {
                 size_t i = 0;
                 for ( ; i < embd.size(); i++) {
-                    // TODO: are the session tokens guaranteed to all be matching here?
-                    //       Should n_matching_session_tokens be re-used instead?
                     if (embd[i] != session_tokens[n_session_consumed]) {
                         session_tokens.resize(n_session_consumed);
                         break;
diff --git a/examples/parallel/parallel.cpp b/examples/parallel/parallel.cpp
@@ -199,7 +199,7 @@ int main(int argc, char ** argv) {
 
         // assign the system KV cache to all parallel sequences
         for (int32_t i = 1; i <= n_clients; ++i) {
-            llama_past_seq_cp(ctx, 0, i, -1, -1);
+            llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);
         }
 
         LOG_INF("\n");
@@ -231,9 +231,9 @@ int main(int argc, char ** argv) {
         if (batch.n_tokens == 0) {
             // all sequences have ended - clear the entire KV cache
             for (int i = 1; i <= n_clients; ++i) {
-                llama_past_seq_rm(ctx, i, -1, -1);
+                llama_kv_cache_seq_rm(ctx, i, -1, -1);
                 // but keep the system prompt
-                llama_past_seq_cp(ctx, 0, i, -1, -1);
+                llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);
             }
 
             LOG_INF("%s: clearing the KV cache\n", __func__);
@@ -370,8 +370,8 @@ int main(int argc, char ** argv) {
                     }
 
                     // delete only the generated part of the sequence, i.e. keep the system prompt in the cache
-                    llama_past_seq_rm(ctx,    client.id + 1, -1, -1);
-                    llama_past_seq_cp(ctx, 0, client.id + 1, -1, -1);
+                    llama_kv_cache_seq_rm(ctx,    client.id + 1, -1, -1);
+                    llama_kv_cache_seq_cp(ctx, 0, client.id + 1, -1, -1);
 
                     const auto t_main_end = ggml_time_us();
 
diff --git a/examples/passkey/passkey.cpp b/examples/passkey/passkey.cpp
@@ -130,11 +130,11 @@ int main(int argc, char ** argv) {
             const int ib = i/n_batch - 1;
             const int bd = n_batch_grp*(n_grp - 1);
 
-            llama_past_seq_add   (ctx, 0, n_past - n_batch,         n_past,         ib*bd);
-            llama_past_seq_div   (ctx, 0, n_past - n_batch + ib*bd, n_past + ib*bd, n_grp);
-            llama_kv_cache_update(ctx);
+            llama_kv_cache_seq_add (ctx, 0, n_past - n_batch,         n_past,         ib*bd);
+            llama_kv_cache_seq_div (ctx, 0, n_past - n_batch + ib*bd, n_past + ib*bd, n_grp);
+            llama_kv_cache_update  (ctx);
 
-            n_past = llama_past_seq_pos_max(ctx, 0) + 1;
+            n_past = llama_kv_cache_seq_pos_max(ctx, 0) + 1;
         }
 
         common_batch_clear(batch);
@@ -164,12 +164,12 @@ int main(int argc, char ** argv) {
 
         LOG_INF("%s: shifting KV cache with %d\n", __func__, n_discard);
 
-        llama_past_seq_rm    (ctx, 0, n_keep            , n_keep + n_discard);
-        llama_past_seq_add   (ctx, 0, n_keep + n_discard, n_ctx,  -n_discard);
-      //llama_kv_cache_defrag(ctx);
-        llama_kv_cache_update(ctx);
+        llama_kv_cache_seq_rm (ctx, 0, n_keep            , n_keep + n_discard);
+        llama_kv_cache_seq_add(ctx, 0, n_keep + n_discard, n_ctx,  -n_discard);
+      //llama_kv_cache_defrag (ctx);
+        llama_kv_cache_update (ctx);
 
-        n_past = llama_past_seq_pos_max(ctx, 0) + 1;
+        n_past = llama_kv_cache_seq_pos_max(ctx, 0) + 1;
 
         common_batch_clear(batch);
 
@@ -195,12 +195,12 @@ int main(int argc, char ** argv) {
         if (n_discard > 0) {
             LOG_INF("%s: shifting KV cache with %d to free space for the answer\n", __func__, n_discard);
 
-            llama_past_seq_rm    (ctx, 0, n_keep            , n_keep + n_discard);
-            llama_past_seq_add   (ctx, 0, n_keep + n_discard, n_ctx,  -n_discard);
-          //llama_kv_cache_defrag(ctx);
-            llama_kv_cache_update(ctx);
+            llama_kv_cache_seq_rm (ctx, 0, n_keep            , n_keep + n_discard);
+            llama_kv_cache_seq_add(ctx, 0, n_keep + n_discard, n_ctx,  -n_discard);
+          //llama_kv_cache_defrag (ctx);
+            llama_kv_cache_update (ctx);
 
-            n_past = llama_past_seq_pos_max(ctx, 0) + 1;
+            n_past = llama_kv_cache_seq_pos_max(ctx, 0) + 1;
         }
     }
 
diff --git a/examples/perplexity/perplexity.cpp b/examples/perplexity/perplexity.cpp
diff --git a/examples/retrieval/retrieval.cpp b/examples/retrieval/retrieval.cpp
diff --git a/examples/save-load-state/save-load-state.cpp b/examples/save-load-state/save-load-state.cpp
diff --git a/examples/server/server.cpp b/examples/server/server.cpp
diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c
diff --git a/include/llama.h b/include/llama.h
diff --git a/src/llama.cpp b/src/llama.cpp

Original file line number	Diff line number	Diff line change
`@@ -966,7 +966,7 @@ struct common_init_result common_init_from_params(common_params & params) {`
`966`	`966`	`if (llama_model_has_decoder(model)) {`
`967`	`967`	`llama_decode(lctx, llama_batch_get_one(tmp.data(), std::min(tmp.size(), (size_t) params.n_batch), 0, 0));`
`968`	`968`	`}`
`969`		`- llama_past_clear(lctx);`
	`969`	`+ llama_kv_cache_clear(lctx);`
`970`	`970`	`llama_synchronize(lctx);`
`971`	`971`	`llama_perf_context_reset(lctx);`
`972`	`972`	`}`
Original file line number	Diff line number	Diff line change
`@@ -133,7 +133,7 @@ int main(int argc, char ** argv) {`
`133`	`133`
`134`	`134`	`const auto t_pp_start = ggml_time_us();`
`135`	`135`
`136`		`- llama_past_clear(ctx);`
	`136`	`+ llama_kv_cache_clear(ctx);`
`137`	`137`
`138`	`138`	`if (!decode_helper(ctx, batch, ctx_params.n_batch)) {`
`139`	`139`	`LOG_ERR("%s: llama_decode() failed\n", __func__);`
`@@ -142,7 +142,7 @@ int main(int argc, char ** argv) {`
`142`	`142`
`143`	`143`	`if (is_pp_shared) {`
`144`	`144`	`for (int32_t i = 1; i < pl; ++i) {`
`145`		`- llama_past_seq_cp(ctx, 0, i, -1, -1);`
	`145`	`+ llama_kv_cache_seq_cp(ctx, 0, i, -1, -1);`
`146`	`146`	`}`
`147`	`147`	`}`
`148`	`148`
Original file line number	Diff line number	Diff line change
`@@ -111,7 +111,7 @@ if llama_decode(context, batch) != 0 {`
`111`	`111`	`}`
`112`	`112`
`113`	`113`	`for i in 1 ..< n_parallel {`
`114`		`- llama_past_seq_cp(context, 0, Int32(i), -1, -1)`
	`114`	`+ llama_kv_cache_seq_cp(context, 0, Int32(i), -1, -1)`
`115`	`115`	`}`
`116`	`116`
`117`	`117`	`if n_parallel > 1 {`
Original file line number	Diff line number	Diff line change
`@@ -338,7 +338,7 @@ static bool cb_eval(struct ggml_tensor * t, bool ask, void * user_data) {`
`338`	`338`	`}`
`339`	`339`
`340`	`340`	`static bool get_hidden_layers(llama_context * ctx, std::vector<llama_token> & tokens) {`
`341`		`- llama_past_clear(ctx);`
	`341`	`+ llama_kv_cache_clear(ctx);`
`342`	`342`	`if (llama_decode(ctx, llama_batch_get_one(tokens.data(), tokens.size(), 0, 0))) {`
`343`	`343`	`fprintf(stderr, "%s : failed to eval\n", __func__);`
`344`	`344`	`return false;`