Enable backend sampling with token speculation

examples/lookup/lookup.cpp

@@ -3,9 +3,11 @@
 #include "common.h"
 #include "ngram-cache.h"
 #include "sampling.h"
+#include "speculative.h"
 #include "log.h"
 #include "llama.h"
 
+#include <algorithm>
 #include <clocale>
 #include <cstdint>
 #include <cstdio>
@@ -27,6 +29,8 @@ int main(int argc, char ** argv){
     // max. number of additional tokens to draft if match is found
     const int n_draft = params.speculative.draft.n_max;
 
+    params.n_outputs_max = common_speculative_n_outputs_max(params.n_batch, params.n_parallel, n_draft);
+
     // init llama.cpp
     llama_backend_init();
     llama_numa_init(params.numa);

examples/speculative-simple/speculative-simple.cpp

@@ -5,6 +5,7 @@
 #include "log.h"
 #include "llama.h"
 
+#include <algorithm>
 #include <clocale>
 #include <cstdio>
 #include <cstring>
@@ -29,6 +30,9 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    params.n_outputs_max = common_speculative_n_outputs_max(
+            params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
+
     // init llama.cpp
     llama_backend_init();
     llama_numa_init(params.numa);

examples/speculative/speculative.cpp

@@ -1,6 +1,7 @@
 #include "arg.h"
 #include "common.h"
 #include "sampling.h"
+#include "speculative.h"
 #include "log.h"
 #include "llama.h"
 
@@ -57,6 +58,9 @@ int main(int argc, char ** argv) {
     // max number of parallel drafting sequences (i.e. tree branches)
     const int n_seq_dft = params.n_parallel;
 
+    params.n_outputs_max = common_speculative_n_outputs_max(
+            params.n_batch, params.n_parallel, params.speculative.draft.n_max);
+
     // probability threshold for splitting a draft branch (only for n_seq_dft > 1)
     const float p_draft_split = params.speculative.draft.p_split;