llama : add llama_max_parallel_sequences()

ggml-ci

Author: ggerganov

include/llama.h

@@ -471,6 +471,7 @@ extern "C" {
     LLAMA_API int64_t llama_time_us(void);
 
     LLAMA_API size_t llama_max_devices(void);
+    LLAMA_API size_t llama_max_parallel_sequences(void);
 
     LLAMA_API bool llama_supports_mmap       (void);
     LLAMA_API bool llama_supports_mlock      (void);

src/llama-context.cpp

@@ -25,7 +25,12 @@ llama_context::llama_context(
 
     const auto & hparams = model.hparams;
 
-    cparams.n_seq_max        = std::max(1u, params.n_seq_max);
+    cparams.n_seq_max = std::max(1u, params.n_seq_max);
+    if (cparams.n_seq_max > LLAMA_MAX_PARALLEL_SEQUENCES) {
+        LLAMA_LOG_WARN("%s: n_seq_max (%d) is larger than the maximum supported (%d) - clamping\n", __func__, cparams.n_seq_max, LLAMA_MAX_PARALLEL_SEQUENCES);
+        cparams.n_seq_max = LLAMA_MAX_PARALLEL_SEQUENCES;
+    }
+
     cparams.n_threads        = params.n_threads;
     cparams.n_threads_batch  = params.n_threads_batch;
     cparams.yarn_ext_factor  = params.yarn_ext_factor;

src/llama-impl.h

@@ -5,6 +5,8 @@
 #include <string>
 #include <vector>
 
+#define LLAMA_MAX_PARALLEL_SEQUENCES 64
+
 #ifdef __GNUC__
 #    if defined(__MINGW32__) && !defined(__clang__)
 #        define LLAMA_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__)))

src/llama-kv-cells.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "llama.h"
+#include "llama-impl.h"
 
 #include <bitset>
 #include <cassert>
@@ -119,7 +119,7 @@ class llama_kv_cells_unified {
         seq[i].reset(seq_id);
 
         if (seq[i].none()) {
-            pos[i]= -1;
+            pos[i] = -1;
 
             used--;
 
@@ -267,6 +267,6 @@ class llama_kv_cells_unified {
     std::vector<llama_pos> shift;
 
     // TODO: assert n_seq_max <= 64
-    std::vector<std::bitset<64>> seq;
+    std::vector<std::bitset<LLAMA_MAX_PARALLEL_SEQUENCES>> seq;
 };
 

src/llama.cpp

@@ -37,6 +37,10 @@ size_t llama_max_devices(void) {
     return 16;
 }
 
+size_t llama_max_parallel_sequences(void) {
+    return LLAMA_MAX_PARALLEL_SEQUENCES;
+}
+
 bool llama_supports_mmap(void) {
     return llama_mmap::SUPPORTED;
 }