llama : move sampling rngs from common to llama

ggml-ci

Author: ggerganov

common/sampling.cpp

@@ -1,11 +1,13 @@
-#define LLAMA_API_INTERNAL
 #include "sampling.h"
+
 #include <random>
 
-struct llama_sampling_context * llama_sampling_init(const struct llama_sampling_params & params) {
+struct llama_sampling_context * llama_sampling_init(const struct llama_sampling_params & params, struct llama_context * ctx, llama_seq_id seq_id) {
     struct llama_sampling_context * result = new llama_sampling_context();
 
     result->params  = params;
+    result->seq_id  = seq_id;
+    result->ctx     = ctx;
     result->grammar = nullptr;
 
     // if there is a grammar, parse it
@@ -81,7 +83,7 @@ void llama_sampling_set_rng_seed(struct llama_sampling_context * ctx, uint32_t s
     if (seed == LLAMA_DEFAULT_SEED) {
         seed = std::random_device{}();
     }
-    ctx->rng.seed(seed);
+    llama_set_rng_seed_seq(ctx->ctx, seed, ctx->seq_id);
 }
 
 void llama_sampling_cp(llama_sampling_context * src, llama_sampling_context * dst) {
@@ -271,10 +273,10 @@ static llama_token llama_sampling_sample_impl(
                   bool is_resampling) {
     const llama_sampling_params & params = ctx_sampling->params;
 
-    const float   temp            = params.temp;
-    const int     mirostat        = params.mirostat;
-    const float   mirostat_tau    = params.mirostat_tau;
-    const float   mirostat_eta    = params.mirostat_eta;
+    const float temp         = params.temp;
+    const int   mirostat     = params.mirostat;
+    const float mirostat_tau = params.mirostat_tau;
+    const float mirostat_eta = params.mirostat_eta;
 
     std::vector<float> original_logits;
     auto cur_p = llama_sampling_prepare(ctx_sampling, ctx_main, ctx_cfg, idx, /* apply_grammar= */ is_resampling, &original_logits);
@@ -304,7 +306,7 @@ static llama_token llama_sampling_sample_impl(
 
             sampler_queue(ctx_main, params, cur_p, min_keep);
 
-            id = llama_sample_token_with_rng(ctx_main, &cur_p, ctx_sampling->rng);
+            id = llama_sample_token_seq(ctx_main, &cur_p, ctx_sampling->seq_id);
 
             //{
             //    const int n_top = 10;

common/sampling.h

@@ -70,9 +70,12 @@ struct llama_sampling_context {
     // parameters that will be used for sampling
     llama_sampling_params params;
 
+    llama_seq_id seq_id;
+
     // mirostat sampler state
     float mirostat_mu;
 
+    llama_context * ctx; // TMP
     llama_grammar * grammar;
 
     // internal
@@ -81,15 +84,14 @@ struct llama_sampling_context {
     // TODO: replace with ring-buffer
     std::vector<llama_token>      prev;
     std::vector<llama_token_data> cur;
-    size_t n_valid; // Number of correct top tokens with correct probabilities.
 
-    std::mt19937 rng;
+    size_t n_valid; // Number of correct top tokens with correct probabilities.
 };
 
 #include "common.h"
 
 // Create a new sampling context instance.
-struct llama_sampling_context * llama_sampling_init(const struct llama_sampling_params & params);
+struct llama_sampling_context * llama_sampling_init(const struct llama_sampling_params & params, struct llama_context * ctx, llama_seq_id seq_id);
 
 void llama_sampling_free(struct llama_sampling_context * ctx);
 

examples/gbnf-validator/gbnf-validator.cpp

@@ -1,8 +1,7 @@
-#define LLAMA_API_INTERNAL
-
 #include "grammar-parser.h"
 #include "ggml.h"
 #include "llama.h"
+#include "llama-impl.h"
 #include "unicode.h"
 
 #include <cstdio>

examples/infill/infill.cpp

@@ -346,7 +346,7 @@ int main(int argc, char ** argv) {
 
     std::vector<llama_token> embd;
 
-    struct llama_sampling_context * ctx_sampling = llama_sampling_init(sparams);
+    struct llama_sampling_context * ctx_sampling = llama_sampling_init(sparams, ctx, 0);
 
     while (n_remain != 0 || params.interactive) {
         // predict

examples/llava/llava-cli.cpp

@@ -191,7 +191,7 @@ static void process_prompt(struct llava_context * ctx_llava, struct llava_image_
 
     LOG_TEE("\n");
 
-    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params->sparams);
+    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params->sparams, ctx_llava->ctx_llama, 0);
     if (!ctx_sampling) {
         fprintf(stderr, "%s: failed to initialize sampling subsystem\n", __func__);
         exit(1);

examples/lookahead/lookahead.cpp

@@ -118,7 +118,7 @@ int main(int argc, char ** argv) {
     llama_batch batch = llama_batch_init(params.n_ctx, 0, W + G + 1);
 
     // target model sampling context
-    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams);
+    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams, ctx, 0);
 
     // verification n-grams
     std::vector<ngram_data> ngrams_cur(G);

examples/lookup/lookup.cpp

@@ -106,7 +106,7 @@ int main(int argc, char ** argv){
 
     bool has_eos = false;
 
-    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams);
+    struct llama_sampling_context * ctx_sampling = llama_sampling_init(params.sparams, ctx, 0);
 
     std::vector<llama_token> draft;
 

examples/main/main.cpp

@@ -527,7 +527,7 @@ int main(int argc, char ** argv) {
         antiprompt_ids.emplace_back(::llama_tokenize(ctx, antiprompt, false, true));
     }
 
-    struct llama_sampling_context * ctx_sampling = llama_sampling_init(sparams);
+    struct llama_sampling_context * ctx_sampling = llama_sampling_init(sparams, ctx, 0);
     if (!ctx_sampling) {
         fprintf(stderr, "%s: failed to initialize sampling subsystem\n", __func__);
         exit(1);

examples/parallel/parallel.cpp

@@ -161,7 +161,7 @@ int main(int argc, char ** argv) {
     for (size_t i = 0; i < clients.size(); ++i) {
         auto & client = clients[i];
         client.id = i;
-        client.ctx_sampling = llama_sampling_init(params.sparams);
+        client.ctx_sampling = llama_sampling_init(params.sparams, ctx, i);
     }
 
     std::vector<llama_token> tokens_system;

examples/quantize-stats/quantize-stats.cpp

@@ -1,7 +1,7 @@
-#define LLAMA_API_INTERNAL
 #include "common.h"
 #include "ggml.h"
 #include "llama.h"
+#include "llama-impl.h"
 
 #include <algorithm>
 #include <cassert>