lookahead : initial working implementation

examples/lookahead/lookahead.cpp

@@ -9,6 +9,10 @@
 struct seq_ngram {
     bool active = false;
 
+    llama_seq_id seq_id = -1;
+
+    std::vector<int> i_batch;
+
     std::vector<llama_token> tokens;
 };
 
@@ -34,9 +38,9 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
-    const int W = 5; // lookahead window
+    const int W = 10; // lookahead window
-    const int N = 4; // n-gram size
+    const int N = 8;  // n-gram size
-    const int G = 5; // max verification n-grams
+    const int G = 10; // max verification n-grams
 
     const bool dump_kv_cache = params.dump_kv_cache;
 
@@ -89,7 +93,7 @@ int main(int argc, char ** argv) {
     llama_decode(ctx, llama_batch_get_one( inp.data(), n_input - 1, 0,           0));
     llama_decode(ctx, llama_batch_get_one(&inp.back(),           1, n_input - 1, 0));
 
-    for (int s = 0; s < W + G + 1; ++s) {
+    for (int s = 1; s < W + G + 1; ++s) {
         llama_kv_cache_seq_cp(ctx, 0, s, -1, -1);
     }
 
@@ -114,15 +118,18 @@ int main(int argc, char ** argv) {
     struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams);
 
     // verification n-grams
-    std::vector<seq_ngram> ngrams(G);
+    std::vector<seq_ngram> ngrams_cur(G);
 
     // tokens for the past N - 1 Jacobi iterations
     std::vector<llama_token> tokens_j_prev(W);
     std::vector<std::vector<llama_token>> tokens_j(N - 1);
     for (int j = 0; j < N - 1; j++) {
         tokens_j[j].resize(W);
+
         for (int i = 0; i < W; i++) {
-            tokens_j[j][i] = all[1 + rand() % (all.size() - 1)];
+            // initialize randomly from the prompt tokens
+            //tokens_j[j][i] = all[1 + rand() % (all.size() - 1)];
+            tokens_j[j][i] = 100 + i;
         }
     }
 
@@ -168,36 +175,95 @@ int main(int argc, char ** argv) {
         {
             llama_batch_clear(batch);
 
+            // current token - first token of the first level
             llama_batch_add(batch, id, n_past, seq_id_all, true);
 
+            // verification n-grams - queue this here for less KV cache fragmentation
+            {
+                const int g_cur = ngrams_observed.cnt[id];
+
+                ngrams_cur.resize(g_cur);
+                for (int g = 0; g < g_cur; g++) {
+                    ngrams_cur[g].active = true;
+                    ngrams_cur[g].tokens.resize(N);
+                    ngrams_cur[g].i_batch.resize(N);
+                    ngrams_cur[g].seq_id = W + 1 + g;
+                    ngrams_cur[g].i_batch[0] = 0;
+                    ngrams_cur[g].tokens [0] = id;
+                }
+
+                for (int j = 0; j < N - 1; j++) {
+                    for (int g = 0; g < g_cur; g++) {
+                        const int idx = id*(N - 1)*G + g*(N - 1);
+
+                        const llama_token t = ngrams_observed.tokens[idx + j];
+
+                        ngrams_cur[g].tokens [j + 1] = t;
+                        ngrams_cur[g].i_batch[j + 1] = batch.n_tokens;
+
+                        llama_batch_add(batch, t, n_past + j + 1, { W + 1 + g }, true);
+                    }
+                }
+            }
+
+            // fill the remaining W - 1 tokens for the first level
             for (int i = 1; i < W; i++) {
                 llama_batch_add(batch, tokens_j[0][i], n_past + i, seq_id_look, false);
             }
 
+            // fill the rest of the levels
             for (int j = 1; j < N - 1; j++) {
                 for (int i = 0; i < W; i++) {
                     llama_batch_add(batch, tokens_j[j][i], n_past + j + i, { i + 1 }, j == N - 2);
                 }
             }
-
-            // TODO: add verification n-grams
         }
 
-        llama_decode(ctx, batch);
+        if (llama_decode(ctx, batch) != 0) {
+            fprintf(stderr, "\n\n%s: error: llama_decode failed - increase KV cache size\n", __func__);
+            return 1;
+        }
 
-        id = llama_sampling_sample(ctx_sampling, ctx, NULL, 0);
+        int seq_id_best = 0;
+
+        for (int v = 0; v < N; ++v) {
+            int i_batch = 0;
+
+            if (v > 0) {
+                for (int g = 0; g < (int) ngrams_cur.size(); g++) {
+                    if (ngrams_cur[g].active) {
+                        i_batch = ngrams_cur[g].i_batch[v];
+                        seq_id_best = ngrams_cur[g].seq_id;
+                        break;
+                    }
+                }
+
+                // no more matches
+                if (i_batch == 0) {
+                    break;
+                }
+            }
+
+            id = llama_sampling_sample(ctx_sampling, ctx, NULL, i_batch);
 
             llama_sampling_accept(ctx_sampling, ctx, id, true);
 
             {
                 const std::string token_str = llama_token_to_piece(ctx, id);
 
+                if (v == 0) {
                     printf("%s", token_str.c_str());
+                } else {
+                    // print light cyan
+                    printf("\033[0;96m%s\033[0m", token_str.c_str());
+                }
                 fflush(stdout);
 
                 if (id == llama_token_eos(model)) {
                     has_eos = true;
                 }
+
+                all.push_back(id);
             }
 
             ++n_predict;
@@ -207,8 +273,22 @@ int main(int argc, char ** argv) {
                 break;
             }
 
+            // verify across active n-grams
+            for (int g = 0; g < (int) ngrams_cur.size(); g++) {
+                if (ngrams_cur[g].active) {
+                    if (v == N - 1) {
+                        ngrams_cur[g].active = false;
+                    } else {
+                        if (id != ngrams_cur[g].tokens[v + 1]) {
+                            ngrams_cur[g].active = false;
+                        } else {
+                        }
+                    }
+                }
+            }
+
             // print known n-grams starting with token id
-        if (1) {
+            if (0) {
                 if (ngrams_observed.cnt[id] > 0) {
                     printf("\n - %d n-grams starting with '%s'\n", ngrams_observed.cnt[id], llama_token_to_piece(ctx, id).c_str());
                 }
@@ -238,8 +318,17 @@ int main(int argc, char ** argv) {
                     tokens_j[j] = tokens_j[j + 1];
                 }
 
+                if (v == 0) {
+                    // sample from the last level
                     for (int i = 0; i < W; i++) {
-                tokens_j[N - 2][i] = llama_sampling_sample(ctx_sampling, ctx, NULL, W*(N - 2) + i);
+                        tokens_j[N - 2][i] = llama_sampling_sample(ctx_sampling, ctx, NULL, ngrams_cur.size()*(N-1) + W*(N - 2) + i);
+                    }
+                } else {
+                    for (int i = 0; i < W; i++) {
+                        // random init
+                        //tokens_j[N - 2][i] = all[1 + rand() % (all.size() - 1)];
+                        tokens_j[N - 2][i] = tokens_j[0][i];
+                    }
                 }
             }
 
@@ -249,6 +338,7 @@ int main(int argc, char ** argv) {
                 std::vector<llama_token> ngram(N - 1);
 
                 // n-gram generation
+                // ref: https://github.com/hao-ai-lab/LookaheadDecoding/issues/14#issuecomment-1826198518
                 for (int f = 0; f < W; ++f) {
                     for (int j = 0; j < N - 1; ++j) {
                         ngram[j] = tokens_j[j][f];
@@ -268,13 +358,19 @@ int main(int argc, char ** argv) {
                     ngrams_observed.n_total++;
                 }
             }
-
-        // verification
-        // TODO
-        {
         }
 
         llama_kv_cache_seq_rm(ctx, -1, n_past, -1);
+
+        if (seq_id_best != 0) {
+            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_kv_cache_seq_cp(ctx, 0, s, -1, -1);
+            }
+        }
     }
 
     auto t_dec_end = ggml_time_us();