Use psirng for sampling

Author: nullspook

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "kompute"]
 	path = ggml/src/kompute
 	url = https://github.com/nomic-ai/kompute.git
+[submodule "libpsirngclient"]
+	path = libpsirngclient
+	url = https://github.com/nullspook/libpsirngclient.git

CMakeLists.txt

@@ -199,3 +199,9 @@ if (LLAMA_BUILD_EXAMPLES)
     add_subdirectory(examples)
     add_subdirectory(pocs)
 endif()
+
+#
+# psirng
+#
+
+add_subdirectory(libpsirngclient)

src/CMakeLists.txt

@@ -22,10 +22,10 @@ add_library(llama
             unicode-data.cpp
             )
 
-target_include_directories(llama PUBLIC . ../include)
+target_include_directories(llama PUBLIC . ../include ../libpsirngclient/src)
 target_compile_features   (llama PUBLIC cxx_std_11) # don't bump
 
-target_link_libraries(llama PUBLIC ggml)
+target_link_libraries(llama PUBLIC ggml psirngclient)
 
 if (BUILD_SHARED_LIBS)
     set_target_properties(llama PROPERTIES POSITION_INDEPENDENT_CODE ON)

src/llama-sampling.cpp

@@ -619,8 +619,21 @@ llama_token llama_sample_token_with_rng_impl(struct llama_sampling * smpl, llama
         probs.push_back(candidates->data[i].p);
     }
 
-    std::discrete_distribution<> dist(probs.begin(), probs.end());
-    int idx = dist(rng);
+    std::vector<float> cdf(probs.size());
+    cdf[0] = probs[0];
+    for (size_t i = 1; i < probs.size(); ++i) {
+        cdf[i] = cdf[i - 1] + probs[i];
+    }
+
+    int idx;
+    double u;
+
+    int rand_result = psirngclient_randuniform(smpl->psirngclient_ptr, &u, 1, 0.0, 1.0);
+    if (rand_result != PSIRNGCLIENT_RESULT_OK) {
+        GGML_ABORT("psirngclient_randuniform error: %d", rand_result);
+    }
+
+    idx = static_cast<int>(std::distance(cdf.begin(), std::lower_bound(cdf.begin(), cdf.end(), u)));
 
     llama_token result = candidates->data[idx].id;
 

src/llama-sampling.h

@@ -2,11 +2,15 @@
 
 #include "llama-impl.h"
 
+#include "psirngclient.h"
+
 struct llama_sampling {
     llama_sampling(int32_t n_vocab) : n_vocab(n_vocab) {}
 
     std::mt19937 rng;
 
+    psirngclient * psirngclient_ptr;
+
     int32_t n_vocab = 0;
 
     mutable int64_t t_sample_us = 0;

src/llama.cpp

@@ -18193,6 +18193,28 @@ struct llama_context * llama_new_context_with_model(
     ctx->abort_callback_data = params.abort_callback_data;
 
     ctx->sampling.rng = std::mt19937(params.seed);
+
+    ctx->sampling.psirngclient_ptr = nullptr;
+
+    const char* psirng_host      = std::getenv("PSIRNG_HOST");
+    const char* psirng_grpc_port = std::getenv("PSIRNG_GRPC_PORT");
+    const char* psirng_cert_path = std::getenv("PSIRNG_CERT_PATH");
+
+    if (psirng_host != nullptr && psirng_grpc_port != nullptr && psirng_cert_path != nullptr) {
+        psirngclient_init(&ctx->sampling.psirngclient_ptr, psirng_host, std::stoi(psirng_grpc_port), psirng_cert_path);
+        if (!psirngclient_ishealthy(ctx->sampling.psirngclient_ptr)) {
+            LLAMA_LOG_ERROR("%s: psirng is not healthy\n", __func__);
+            llama_free(ctx);
+            return nullptr;
+        } else {
+            LLAMA_LOG_INFO("%s: Using psirng running on %s:%s\n", __func__, psirng_host, psirng_grpc_port);
+        }
+    } else {
+        LLAMA_LOG_ERROR("%s: psirng is not configured\n", __func__);
+        llama_free(ctx);
+        return nullptr;
+    }
+
     ctx->logits_all   = params.logits_all;
     // build worst-case graph for encoder if a model contains encoder
     ctx->is_encoding  = llama_model_has_encoder(model);