diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 690f768ca4..6cf5971683 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -1114,7 +1114,11 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da cur_p->selected = 0; + std::uniform_real_distribution dist(0.0f, 1.0f); + if (cur_p->size == 1) { + // keep the RNG state aligned with backend sampling, which draws once per output + dist(ctx->rng); cur_p->data[0].p = 1.0f; return; } @@ -1139,7 +1143,6 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da // sample from the obtained probabilities and normalize the probs in a single pass // this is ~3x faster on Mac with full gpt-oss vocab than the version below // - std::uniform_real_distribution dist(0.0f, 1.0f); const double rnd = dist(ctx->rng); double sum_run = 0.0f; diff --git a/tests/test-sampling.cpp b/tests/test-sampling.cpp index df1eb1a202..7f22040f3d 100644 --- a/tests/test-sampling.cpp +++ b/tests/test-sampling.cpp @@ -61,6 +61,31 @@ private: std::vector cur; }; +static llama_token sample_dist(llama_sampler * sampler, const std::vector & logits) { + std::vector cur; + for (llama_token token_id = 0; token_id < (llama_token) logits.size(); ++token_id) { + cur.push_back({ token_id, logits[token_id], 0.0f }); + } + + llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false }; + llama_sampler_apply(sampler, &cur_p); + GGML_ASSERT(cur_p.selected >= 0); + return cur_p.data[cur_p.selected].id; +} + +static void test_dist_singleton_rng() { + llama_sampler * singleton = llama_sampler_init_dist(4242); + llama_sampler * control = llama_sampler_init_dist(4242); + + sample_dist(singleton, { 0.0f }); + sample_dist(control, { 0.0f, 0.0f }); + + GGML_ASSERT(sample_dist(singleton, { 0.0f, 0.0f }) == sample_dist(control, { 0.0f, 0.0f })); + + llama_sampler_free(singleton); + llama_sampler_free(control); +} + static void test_temp(const std::vector & probs, const std::vector & probs_expected, float temp) { sampler_tester tester(probs, probs_expected); @@ -308,6 +333,8 @@ static void test_perf() { int main(void) { ggml_time_init(); + test_dist_singleton_rng(); + test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 1.0f); test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.0f, 0.0f, 0.0f, 1.0f}, 0.0f);