Felipe97/llama-cpp-compiled
01.1k
1#include "llama-sampler.h"2 3#include "llama-impl.h"4#include "llama-vocab.h"5#include "llama-grammar.h"6 7#include "ggml-cpp.h"8 9#include <array>10#include <algorithm>11#include <cassert>12#include <cfloat>13#include <chrono>14#include <cmath>15#include <cstdlib>16#include <cstring>17#include <ctime>18#include <numeric>19#include <random>20#include <unordered_map>21#include <stdexcept>22 23// the ring buffer works similarly to std::deque, but with a fixed capacity24template<typename T>25struct ring_buffer {26 ring_buffer(size_t cap) : capacity(cap), data(cap) {}27 28 T & front() {29 if (sz == 0) {30 throw std::runtime_error("ring buffer is empty");31 }32 return data[first];33 }34 35 const T & front() const {36 if (sz == 0) {37 throw std::runtime_error("ring buffer is empty");38 }39 return data[first];40 }41 42 T & back() {43 if (sz == 0) {44 throw std::runtime_error("ring buffer is empty");45 }46 return data[pos];47 }48 49 const T & back() const {50 if (sz == 0) {51 throw std::runtime_error("ring buffer is empty");52 }53 return data[pos];54 }55 56 void push_back(const T & value) {57 if (capacity == 0) {58 throw std::runtime_error("ring buffer: capacity is zero");59 }60 61 if (sz == capacity) {62 // advance the start when buffer is full63 first = (first + 1) % capacity;64 } else {65 sz++;66 }67 data[pos] = value;68 pos = (pos + 1) % capacity;69 }70 71 T pop_front() {72 if (sz == 0) {73 throw std::runtime_error("ring buffer is empty");74 }75 T value = data[first];76 first = (first + 1) % capacity;77 sz--;78 return value;79 }80 81 //T & operator[](size_t i) {82 // if (i >= sz) {83 // throw std::runtime_error("ring buffer: index out of bounds");84 // }85 // return data[(first + i) % capacity];86 //}87 88 //const T & at(size_t i) const {89 // if (i >= sz) {90 // throw std::runtime_error("ring buffer: index out of bounds");91 // }92 // return data[(first + i) % capacity];93 //}94 95 const T & rat(size_t i) const {96 if (i >= sz) {97 throw std::runtime_error("ring buffer: index out of bounds");98 }99 return data[(first + sz - i - 1) % capacity];100 }101 102 std::vector<T> to_vector() const {103 std::vector<T> result;104 result.reserve(sz);105 for (size_t i = 0; i < sz; i++) {106 result.push_back(data[(first + i) % capacity]);107 }108 return result;109 }110 111 void clear() {112 // here only reset the status of the buffer113 sz = 0;114 first = 0;115 pos = 0;116 }117 118 bool empty() const {119 return sz == 0;120 }121 122 size_t size() const {123 return sz;124 }125 126 size_t capacity = 0;127 size_t sz = 0;128 size_t first = 0;129 size_t pos = 0;130 131 std::vector<T> data;132};133 134// writes result in res, does not mutate cur135static void llama_token_data_array_partial_sort(const llama_token_data_array & cur, int npartial, std::vector<llama_token_data> & res) {136 static const auto comp = [](const llama_token_data & a, const llama_token_data & b) {137 return a.logit > b.logit;138 };139 140 constexpr int nbuckets = 128;141 constexpr float bucket_low = -10.0f;142 constexpr float bucket_high = 10.0f;143 constexpr float bucket_scale = nbuckets/(bucket_high - bucket_low);144 constexpr float bucket_inter = -bucket_low * bucket_scale;145 146 std::vector<int> bucket_idx;147 std::vector<int> histo(nbuckets, 0);148 149 std::vector<llama_token_data*> bucket_ptrs;150 151 bucket_idx.reserve(cur.size);152 153 for (int i = 0; i < (int)cur.size; ++i) {154 const float val = cur.data[i].logit;155 int ib = int(bucket_scale * val + bucket_inter); //nbuckets * (val - bucket_low) / (bucket_high - bucket_low);156 ib = std::max(0, std::min(nbuckets - 1, ib));157 bucket_idx.push_back(ib);158 ++histo[ib];159 }160 int nhave = 0;161 int ib = nbuckets - 1;162 for ( ; ib >= 0; --ib) {163 nhave += histo[ib];164 if (nhave >= npartial) {165 break;166 }167 }168 res.resize(nhave);169 auto * ptr = res.data();170 bucket_ptrs.reserve(nbuckets - ib);171 for (int j = nbuckets - 1; j >= ib; --j) {172 bucket_ptrs.push_back(ptr);173 ptr += histo[j];174 }175 for (int i = 0; i < (int)cur.size; ++i) {176 int j = bucket_idx[i];177 if (j >= ib) {178 *bucket_ptrs[nbuckets - 1 - j]++ = cur.data[i];179 }180 }181 182 ptr = res.data();183 int ndone = 0;184 for (int j = nbuckets - 1; j > ib; --j) {185 std::sort(ptr, ptr + histo[j], comp);186 ptr += histo[j];187 ndone += histo[j];188 }189 std::partial_sort(ptr, ptr + npartial - ndone, ptr + histo[ib], comp);190}191 192// reduces the size of cur_p to npartial, keeping only the top npartial elements193static void llama_token_data_array_partial_sort_inplace(llama_token_data_array * cur_p, int npartial) {194 static const auto comp = [](const llama_token_data & a, const llama_token_data & b) {195 return a.logit > b.logit;196 };197 198 if (npartial <= 128) {199 std::partial_sort(cur_p->data, cur_p->data + npartial, cur_p->data + cur_p->size, comp);200 201 cur_p->size = npartial;202 cur_p->sorted = true;203 204 return;205 }206 207 std::vector<llama_token_data> tmp;208 209 llama_token_data_array_partial_sort(*cur_p, npartial, tmp);210 211 std::copy(tmp.data(), tmp.data() + npartial, cur_p->data);212 213 cur_p->size = npartial;214 cur_p->sorted = true;215}216 217static int llama_sample_dist(llama_token_data_array * cur_p, std::mt19937 & rng) {218 // iterator for the probabilities219#ifdef __GNUC__220 #pragma GCC diagnostic push221 #pragma GCC diagnostic ignored "-Wunused-local-typedefs"222#endif223 224 struct probs_iterator {225 typedef std::input_iterator_tag iterator_category;226 typedef float value_type;227 typedef float * pointer;228 typedef float & reference;229 typedef ptrdiff_t difference_type;230 231 const llama_token_data * data;232 233 bool operator==(const probs_iterator & other) const { return data == other.data; }234 bool operator!=(const probs_iterator & other) const { return data != other.data; }235 const float & operator*() const { return data->p; }236 probs_iterator & operator++() { ++data; return *this; }237 probs_iterator operator++(int) { probs_iterator tmp = *this; ++data; return tmp; }238 };239 240#ifdef __GNUC__241 #pragma GCC diagnostic pop242#endif243 244 std::discrete_distribution<int> dist(probs_iterator{cur_p->data}, probs_iterator{cur_p->data + cur_p->size});245 246 return dist(rng);247}248 249/*250static void llama_log_softmax(float * array, size_t size) {251 float max_l = *std::max_element(array, array + size);252 float sum = 0.f;253 for (size_t i = 0; i < size; ++i) {254 float p = expf(array[i] - max_l);255 sum += p;256 array[i] = p;257 }258 259 for (size_t i = 0; i < size; ++i) {260 array[i] = logf(array[i] / sum);261 }262}263*/264 265static void llama_sampler_temp_impl(llama_token_data_array * cur_p, float temp) {266 if (cur_p->size == 0) {267 return;268 }269 270 if (temp <= 0.0f) {271 // find the token with the highest logit and set the rest to -inf272 size_t max_i = 0;273 float max_l = cur_p->data[0].logit;274 275 for (size_t i = 1; i < cur_p->size; ++i) {276 if (cur_p->data[i ].logit > max_l) {277 cur_p->data[max_i].logit = -INFINITY;278 max_i = i;279 max_l = cur_p->data[i].logit;280 } else {281 cur_p->data[i].logit = -INFINITY;282 }283 }284 285 return;286 }287 288 for (size_t i = 0; i < cur_p->size; ++i) {289 cur_p->data[i].logit /= temp;290 }291}292 293static void llama_sampler_softmax_impl(llama_token_data_array * cur_p, bool do_sort) {294 GGML_ASSERT(cur_p->size > 0);295 296 // Sort the logits in descending order if requested297 if (do_sort && !cur_p->sorted) {298 llama_token_data_array_partial_sort_inplace(cur_p, cur_p->size);299 }300 301 float max_l = cur_p->data[0].logit;302 if (!cur_p->sorted) {303 for (size_t i = 1; i < cur_p->size; ++i) {304 max_l = std::max(max_l, cur_p->data[i].logit);305 }306 }307 308 float cum_sum = 0.0f;309 310 for (size_t i = 0; i < cur_p->size; ++i) {311 float p = expf(cur_p->data[i].logit - max_l);312 cur_p->data[i].p = p;313 cum_sum += p;314 }315 316 for (size_t i = 0; i < cur_p->size; ++i) {317 cur_p->data[i].p /= cum_sum;318 }319}320 321static void llama_sampler_top_k_impl(llama_token_data_array * cur_p, int32_t k) {322 // if (k >= (int32_t)cur_p->size) {323 // return;324 // }325 326 if (k <= 0) {327 return;328 }329 330 k = std::min(k, (int) cur_p->size);331 332 // Sort scores in descending order333 if (!cur_p->sorted) {334 llama_token_data_array_partial_sort_inplace(cur_p, k);335 }336 337 cur_p->size = k;338}339 340static uint32_t get_rng_seed(uint32_t seed) {341 if (seed == LLAMA_DEFAULT_SEED) {342 // use system clock if std::random_device is not a true RNG343 static bool is_rd_prng = std::random_device().entropy() == 0;344 if (is_rd_prng) {345 return (uint32_t) std::chrono::system_clock::now().time_since_epoch().count();346 }347 std::random_device rd;348 return rd();349 }350 return seed;351}352 353// llama_sampler API354 355struct llama_sampler * llama_sampler_init(356 struct llama_sampler_i * iface,357 llama_sampler_context_t ctx) {358 return new llama_sampler {359 /* .iface = */ iface,360 /* .ctx = */ ctx,361 };362}363 364const char * llama_sampler_name(const struct llama_sampler * smpl) {365 if (!smpl->iface) {366 return "(null)";367 }368 369 return smpl->iface->name(smpl);370}371 372void llama_sampler_accept(struct llama_sampler * smpl, llama_token token) {373 if (!smpl) {374 return;375 }376 377 if (smpl->iface->accept) {378 smpl->iface->accept(smpl, token);379 }380}381 382void llama_sampler_apply(struct llama_sampler * smpl, struct llama_token_data_array * cur_p) {383 if (!smpl) {384 return;385 }386 387 GGML_ASSERT(smpl->iface->apply);388 smpl->iface->apply(smpl, cur_p);389}390 391void llama_sampler_reset(struct llama_sampler * smpl) {392 if (!smpl) {393 return;394 }395 396 if (smpl->iface->reset) {397 smpl->iface->reset(smpl);398 }399}400 401struct llama_sampler * llama_sampler_clone(const struct llama_sampler * smpl) {402 if (!smpl) {403 return nullptr;404 }405 406 if (smpl->iface->clone) {407 return smpl->iface->clone(smpl);408 }409 410 if (smpl->ctx == nullptr) {411 return llama_sampler_init(412 /* .iface = */ smpl->iface,413 /* .ctx = */ nullptr414 );415 }416 417 GGML_ABORT("the sampler does not support cloning");418}419 420void llama_sampler_free(struct llama_sampler * smpl) {421 if (smpl == nullptr) {422 return;423 }424 425 if (smpl->iface->free) {426 smpl->iface->free(smpl);427 }428 429 delete smpl;430}431 432// empty sampler433 434struct llama_sampler_empty {435 const char * name;436};437 438static struct llama_sampler * llama_sampler_init_empty(const char * name);439 440static const char * llama_sampler_empty_name(const struct llama_sampler * smpl) {441 auto * ctx = (llama_sampler_empty *) smpl->ctx;442 return ctx->name;443}444 445static void llama_sampler_empty_accept(struct llama_sampler * smpl, llama_token token) {446 GGML_UNUSED(smpl);447 GGML_UNUSED(token);448}449 450static void llama_sampler_empty_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {451 GGML_UNUSED(smpl);452 GGML_UNUSED(cur_p);453}454 455static void llama_sampler_empty_reset(struct llama_sampler * smpl) {456 GGML_UNUSED(smpl);457}458 459static struct llama_sampler * llama_sampler_empty_clone(const struct llama_sampler * smpl) {460 auto * ctx = (llama_sampler_empty *) smpl->ctx;461 return llama_sampler_init_empty(ctx->name);462}463 464static void llama_sampler_empty_free(struct llama_sampler * smpl) {465 delete (llama_sampler_empty *) smpl->ctx;466}467 468static bool llama_sampler_empty_backend_init(469 struct llama_sampler * smpl,470 ggml_backend_buffer_type_t buft,471 uint32_t n_outputs_max_per_seq) {472 GGML_UNUSED(smpl);473 GGML_UNUSED(buft);474 GGML_UNUSED(n_outputs_max_per_seq);475 476 return true;477}478 479static void llama_sampler_empty_backend_accept(480 struct llama_sampler * smpl,481 ggml_context * ctx,482 ggml_cgraph * gf,483 struct ggml_tensor * selected_token) {484 GGML_UNUSED(smpl);485 GGML_UNUSED(ctx);486 GGML_UNUSED(gf);487 GGML_UNUSED(selected_token);488}489 490static void llama_sampler_empty_backend_apply(491 struct llama_sampler * smpl,492 struct ggml_context * ctx,493 struct ggml_cgraph * gf,494 struct llama_sampler_data * data) {495 GGML_UNUSED(smpl);496 GGML_UNUSED(ctx);497 GGML_UNUSED(gf);498 GGML_UNUSED(data);499}500 501static void llama_sampler_empty_backend_set_input(struct llama_sampler * smpl) {502 GGML_UNUSED(smpl);503}504 505static struct llama_sampler_i llama_sampler_empty_i = {506 /* .name = */ llama_sampler_empty_name,507 /* .accept = */ llama_sampler_empty_accept,508 /* .apply = */ llama_sampler_empty_apply,509 /* .reset = */ llama_sampler_empty_reset,510 /* .clone = */ llama_sampler_empty_clone,511 /* .free = */ llama_sampler_empty_free,512 /* .backend_init = */ llama_sampler_empty_backend_init,513 /* .backend_accept = */ llama_sampler_empty_backend_accept,514 /* .backend_apply = */ llama_sampler_empty_backend_apply,515 /* .backend_set_input = */ llama_sampler_empty_backend_set_input,516 /* .backend_reset = */ nullptr,517 /* .copy_state = */ nullptr,518};519 520struct llama_sampler * llama_sampler_init_empty(const char * name) {521 return llama_sampler_init(522 /* .iface = */ &llama_sampler_empty_i,523 /* .ctx = */ new llama_sampler_empty {524 /* .name = */ name,525 }526 );527}528 529// common backend sampler functionality530//531// +name : means that the sampler is support and will run on the backend532// -name : means that a ggml operator is not supported by the backend533//534struct llama_sampler_backend {535 llama_sampler_backend(const char * name) : name(name), name_ext(name), is_init(false), support(false) {}536 537 const char * get_name() {538 if (!is_init) {539 return name.c_str();540 }541 542 if (support) {543 name_ext = "+" + name;544 } else {545 name_ext = "-" + name;546 }547 548 return name_ext.c_str();549 }550 551 void init(bool support) {552 GGML_ASSERT(this->is_init == false);553 554 this->is_init = true;555 this->support = support;556 }557 558 // copy the state that is not tied to the current sampling graph559 // samplers that hold only immutable configuration can use this as is560 void copy_state(const llama_sampler_backend & src) {561 GGML_UNUSED(src);562 }563 564private:565 std::string name;566 std::string name_ext;567 568 bool is_init;569 bool support;570};571 572// .copy_state for samplers deriving from llama_sampler_backend573template<typename T>574static void llama_sampler_backend_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {575 ((T *) dst->ctx)->copy_state(*(const T *) src->ctx);576}577 578struct llama_sampler_backend_probe {579 ggml_context_ptr ctx;580 ggml_cgraph * gf;581};582 583static llama_sampler_backend_probe llama_sampler_backend_probe_graph(584 llama_sampler * sampler,585 int64_t n_candidates,586 uint32_t max_nodes,587 bool with_candidates) {588 ggml_init_params params = {589 /*.mem_size =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),590 /*.mem_buffer =*/ nullptr,591 /*.no_alloc =*/ true,592 };593 594 ggml_context_ptr ctx_ptr { ggml_init(params) };595 if (!ctx_ptr) {596 throw std::runtime_error(format("failed to create ggml context"));597 }598 599 auto * ctx = ctx_ptr.get();600 auto * gf = ggml_new_graph_custom(ctx, max_nodes, false);601 602 llama_sampler_data data = {603 /*.logits =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates),604 /*.probs =*/ nullptr,605 /*.sampled =*/ nullptr,606 /*.candidates =*/ with_candidates ? ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates) : nullptr,607 };608 609 if (sampler->iface->backend_reset) {610 sampler->iface->backend_reset(sampler);611 }612 sampler->iface->backend_apply(sampler, ctx, gf, &data);613 614 for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {615 if (output) {616 ggml_build_forward_expand(gf, output);617 }618 }619 620 if (sampler->iface->backend_reset) {621 sampler->iface->backend_reset(sampler);622 }623 624 return { std::move(ctx_ptr), gf };625}626 627static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) {628 uint32_t n_tensors = 0;629 for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor;630 tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) {631 ++n_tensors;632 }633 634 return std::max<uint32_t>(ggml_graph_n_nodes(probe.gf), n_tensors);635}636 637// check if all ggml ops used by the sampler are supported by the backend638static bool llama_sampler_backend_support(639 llama_sampler * smpl,640 ggml_backend_buffer_type_t buft) {641 auto * device = ggml_backend_buft_get_device(buft);642 if (!device) {643 // CPU backend always supported644 return true;645 }646 647 auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true);648 649 for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) {650 struct ggml_tensor * op = ggml_graph_node(probe.gf, i);651 652 if (!ggml_backend_dev_supports_op(device, op)) {653 LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",654 __func__, ggml_backend_dev_name(device), ggml_op_name(op->op), smpl->iface->name(smpl));655 656 return false;657 }658 }659 660 return true;661}662 663// sampler chain664 665static const char * llama_sampler_chain_name(const struct llama_sampler * /*smpl*/) {666 return "chain";667}668 669static void llama_sampler_chain_accept(struct llama_sampler * smpl, llama_token token) {670 auto * chain = (llama_sampler_chain *) smpl->ctx;671 672 time_meas tm(chain->t_sample_us, chain->params.no_perf);673 674 for (auto & smpl : chain->samplers) {675 llama_sampler_accept(smpl.ptr, token);676 }677 678 chain->n_sample++;679}680 681static void llama_sampler_chain_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {682 auto * chain = (llama_sampler_chain *) smpl->ctx;683 684 time_meas tm(chain->t_sample_us, chain->params.no_perf);685 686 bool is_backend = chain->is_init;687 688 for (auto & smpl : chain->samplers) {689 if (is_backend && smpl.is_backend) {690 continue;691 }692 693 is_backend = false;694 695 if (smpl.ptr->iface->apply == nullptr) {696 continue;697 }698 699 llama_sampler_apply(smpl.ptr, cur_p);700 }701}702 703static void llama_sampler_chain_reset(struct llama_sampler * smpl) {704 auto * chain = (llama_sampler_chain *) smpl->ctx;705 706 for (auto & smpl : chain->samplers) {707 llama_sampler_reset(smpl.ptr);708 }709}710 711static struct llama_sampler * llama_sampler_chain_clone(const struct llama_sampler * smpl) {712 const auto * chain_src = (const llama_sampler_chain *) smpl->ctx;713 714 auto * result = llama_sampler_chain_init(chain_src->params);715 716 for (const auto & smpl : chain_src->samplers) {717 llama_sampler_chain_add(result, llama_sampler_clone(smpl.ptr));718 }719 720 return result;721}722 723static void llama_sampler_chain_free(struct llama_sampler * smpl) {724 auto * chain = (llama_sampler_chain *) smpl->ctx;725 726 for (auto & smpl : chain->samplers) {727 llama_sampler_free(smpl.ptr);728 }729 730 delete chain;731}732 733static bool llama_sampler_chain_backend_init(734 struct llama_sampler * smpl,735 ggml_backend_buffer_type_t buft,736 uint32_t n_outputs_max_per_seq) {737 auto * chain = (llama_sampler_chain *) smpl->ctx;738 739 GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");740 741 chain->is_init = true;742 743 bool res = true;744 bool backend_prefix = true;745 746 for (auto & smpl : chain->samplers) {747 bool cur_prefix = backend_prefix;748 749 // to be able to run a sampler on the backend, it has to:750 // - have the .backend_init() API implemented751 // - return true during .backend_init()752 // - support the requested per-sequence output limit753 if (cur_prefix && smpl.ptr->iface->backend_init) {754 if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) {755 cur_prefix = false;756 }757 } else {758 cur_prefix = false;759 }760 761 smpl.is_backend = cur_prefix;762 backend_prefix = cur_prefix;763 764 res = res && cur_prefix;765 }766 767 auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false);768 chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe);769 770 return res;771}772 773static void llama_sampler_chain_backend_accept(774 struct llama_sampler * smpl,775 ggml_context * ctx,776 ggml_cgraph * gf,777 struct ggml_tensor * selected_token) {778 auto * chain = (llama_sampler_chain *) smpl->ctx;779 780 for (auto & smpl : chain->samplers) {781 if (!smpl.is_backend) {782 break;783 }784 785 if (smpl.ptr->iface->backend_accept) {786 smpl.ptr->iface->backend_accept(smpl.ptr, ctx, gf, selected_token);787 }788 }789}790 791static void llama_sampler_chain_backend_apply(792 struct llama_sampler * smpl,793 struct ggml_context * ctx,794 struct ggml_cgraph * gf,795 struct llama_sampler_data * data) {796 auto * chain = (llama_sampler_chain *) smpl->ctx;797 798 GGML_ASSERT(chain->is_init && "llama_sampler_chain_backend_init() not called");799 800 for (auto & smpl : chain->samplers) {801 if (!smpl.is_backend) {802 break;803 }804 805 if (smpl.ptr->iface->backend_apply) {806 smpl.ptr->iface->backend_apply(smpl.ptr, ctx, gf, data);807 }808 }809}810 811static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {812 auto * chain = (llama_sampler_chain *) smpl->ctx;813 814 for (auto & smpl : chain->samplers) {815 if (!smpl.is_backend) {816 break;817 }818 819 if (smpl.ptr->iface->backend_set_input) {820 smpl.ptr->iface->backend_set_input(smpl.ptr);821 }822 }823}824 825static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) {826 auto * chain = (llama_sampler_chain *) smpl->ctx;827 828 for (auto & entry : chain->samplers) {829 if (!entry.is_backend) {830 break;831 }832 if (entry.ptr->iface->backend_reset) {833 entry.ptr->iface->backend_reset(entry.ptr);834 }835 }836}837 838static void llama_sampler_chain_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {839 const auto * src_chain = (const llama_sampler_chain *) src->ctx;840 auto * dst_chain = (llama_sampler_chain *) dst->ctx;841 842 GGML_ASSERT(src_chain->samplers.size() == dst_chain->samplers.size());843 844 for (size_t i = 0; i < src_chain->samplers.size(); ++i) {845 llama_sampler_copy(src_chain->samplers[i].ptr, dst_chain->samplers[i].ptr);846 }847 848 // note: is_init, n_nodes and is_backend belong to the current sampling graph849 dst_chain->params = src_chain->params;850 dst_chain->cur = src_chain->cur;851 dst_chain->t_sample_us = src_chain->t_sample_us;852 dst_chain->n_sample = src_chain->n_sample;853}854 855static struct llama_sampler_i llama_sampler_chain_i = {856 /* .name = */ llama_sampler_chain_name,857 /* .accept = */ llama_sampler_chain_accept,858 /* .apply = */ llama_sampler_chain_apply,859 /* .reset = */ llama_sampler_chain_reset,860 /* .clone = */ llama_sampler_chain_clone,861 /* .free = */ llama_sampler_chain_free,862 /* .backend_init = */ llama_sampler_chain_backend_init,863 /* .backend_accept = */ llama_sampler_chain_backend_accept,864 /* .backend_apply = */ llama_sampler_chain_backend_apply,865 /* .backend_set_input = */ llama_sampler_chain_backend_set_input,866 /* .backend_reset = */ llama_sampler_chain_backend_reset,867 /* .copy_state = */ llama_sampler_chain_copy_state,868};869 870struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {871 return llama_sampler_init(872 /* .iface = */ &llama_sampler_chain_i,873 /* .ctx = */ new llama_sampler_chain {874 /* .params = */ params,875 /* .is_init = */ false,876 /* .n_nodes = */ 0,877 /* .samplers = */ {},878 /* .cur = */ {},879 /* .t_sample_us = */ 0,880 /* .n_sample = */ 0,881 }882 );883}884 885uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) {886 GGML_ASSERT(sampler != nullptr);887 GGML_ASSERT(sampler->iface == &llama_sampler_chain_i);888 889 const auto * chain = (const llama_sampler_chain *) sampler->ctx;890 GGML_ASSERT(chain->is_init);891 892 return chain->n_nodes;893}894 895llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {896 const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx);897 const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx);898 const float * sampled_logits = llama_get_sampled_logits_ith (ctx, idx);899 const llama_token * sampled_ids = llama_get_sampled_candidates_ith(ctx, idx);900 901 // If a backend sampler has already sampled a token, return it.902 if (sampled_token != LLAMA_TOKEN_NULL) {903 LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);904 llama_sampler_accept(smpl, sampled_token);905 return sampled_token;906 }907 908 const llama_model * model = llama_get_model(ctx);909 const llama_vocab * vocab = llama_model_get_vocab(model);910 911 const int n_vocab = llama_vocab_n_tokens(vocab);912 913 // use pre-allocated buffer from chain if available, otherwise allocate locally914 std::vector<llama_token_data> * cur_ptr;915 std::vector<llama_token_data> cur_local;916 917 if (smpl->iface == &llama_sampler_chain_i) {918 auto * chain = (llama_sampler_chain *) smpl->ctx;919 cur_ptr = &chain->cur;920 } else {921 cur_ptr = &cur_local;922 }923 924 auto & cur = *cur_ptr;925 926 if (sampled_probs) {927 const uint32_t sampled_probs_count = llama_get_sampled_probs_count_ith(ctx, idx);928 cur.resize(sampled_probs_count);929 for (uint32_t i = 0; i < sampled_probs_count; ++i) {930 cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], sampled_probs[i]};931 }932 } else if (sampled_logits) {933 const uint32_t sampled_logits_count = llama_get_sampled_logits_count_ith(ctx, idx);934 cur.resize(sampled_logits_count);935 for (llama_token i = 0; i < (int)sampled_logits_count; i++) {936 cur[i] = llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f};937 }938 } else {939 const auto * logits = llama_get_logits_ith(ctx, idx);940 GGML_ASSERT(logits != nullptr);941 cur.resize(n_vocab);942 for (llama_token token_id = 0; token_id < n_vocab; token_id++) {943 cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f};944 }945 }946 947 llama_token_data_array cur_p = {948 /* .data = */ cur.data(),949 /* .size = */ cur.size(),950 /* .selected = */ -1,951 /* .sorted = */ false,952 };953 954 llama_sampler_apply(smpl, &cur_p);955 956 GGML_ASSERT(cur_p.selected >= 0 && cur_p.selected < (int32_t) cur_p.size);957 958 auto token = cur_p.data[cur_p.selected].id;959 960 llama_sampler_accept(smpl, token);961 962 return token;963}964 965 966void llama_sampler_chain_add(struct llama_sampler * chain, struct llama_sampler * smpl) {967 auto * p = (llama_sampler_chain *) chain->ctx;968 p->samplers.push_back({969 /* .is_backend = */ false,970 /* .ptr = */ smpl,971 });972}973 974struct llama_sampler * llama_sampler_chain_get(struct llama_sampler * chain, int32_t i) {975 if (chain == nullptr) {976 return nullptr;977 }978 979 if (chain->iface != &llama_sampler_chain_i) {980 return nullptr;981 }982 983 if (i == -1) {984 return chain;985 }986 987 const auto * p = (const llama_sampler_chain *) chain->ctx;988 989 if (i < 0 || (size_t) i >= p->samplers.size()) {990 return nullptr;991 }992 993 return p->samplers[i].ptr;994}995 996struct llama_sampler * llama_sampler_chain_remove(struct llama_sampler * chain, int32_t i) {997 auto * p = (llama_sampler_chain *) chain->ctx;998 999 if (i < 0 || (size_t) i >= p->samplers.size()) {1000 return nullptr;1001 }1002 1003 auto * result = p->samplers[i].ptr;1004 p->samplers.erase(p->samplers.begin() + i);1005 1006 return result;1007}1008 1009int32_t llama_sampler_chain_n(const struct llama_sampler * chain) {1010 const auto * p = (const llama_sampler_chain *) chain->ctx;1011 1012 return p->samplers.size();1013}1014 1015//1016// samplers1017//1018 1019// greedy1020 1021struct llama_sampler_greedy : public llama_sampler_backend {1022};1023 1024static const char * llama_sampler_greedy_name(const struct llama_sampler * smpl) {1025 auto * sctx = (llama_sampler_greedy *) smpl->ctx;1026 return sctx->get_name();1027}1028 1029static void llama_sampler_greedy_reset(struct llama_sampler * smpl) {1030 auto * ctx = (llama_sampler_greedy *) smpl->ctx;1031 GGML_UNUSED(ctx);1032}1033 1034static struct llama_sampler * llama_sampler_greedy_clone(const struct llama_sampler * smpl) {1035 const auto * ctx = (const llama_sampler_greedy *) smpl->ctx;1036 auto * result = llama_sampler_init_greedy();1037 1038 // copy the state1039 {1040 auto * result_ctx = (llama_sampler_greedy *) result->ctx;1041 1042 GGML_UNUSED(ctx);1043 GGML_UNUSED(result_ctx);1044 }1045 1046 return result;1047}1048 1049static void llama_sampler_greedy_free(struct llama_sampler * smpl) {1050 delete (llama_sampler_greedy *) smpl->ctx;1051}1052 1053static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_token_data_array * cur_p) {1054 cur_p->selected = 0;1055 for (size_t i = 1; i < cur_p->size; ++i) {1056 if (cur_p->data[i].logit > cur_p->data[cur_p->selected].logit) {1057 cur_p->selected = i;1058 }1059 }1060}1061 1062static bool llama_sampler_greedy_backend_init(1063 struct llama_sampler * smpl,1064 ggml_backend_buffer_type_t buft,1065 uint32_t n_outputs_max_per_seq) {1066 auto * sctx = (llama_sampler_greedy *) smpl->ctx;1067 GGML_UNUSED(n_outputs_max_per_seq);1068 1069 const bool res = llama_sampler_backend_support(smpl, buft);1070 1071 sctx->init(res);1072 1073 return res;1074}1075 1076static void llama_sampler_greedy_backend_apply(1077 struct llama_sampler * smpl,1078 struct ggml_context * ctx,1079 struct ggml_cgraph * gf,1080 struct llama_sampler_data * data) {1081 GGML_UNUSED(gf);1082 GGML_UNUSED(smpl);1083 1084 struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));1085 1086 struct ggml_tensor * curl = ggml_argmax(ctx, logits);1087 ggml_set_name(curl, "greedy_argmax");1088 1089 data->sampled = curl;1090}1091 1092static struct llama_sampler_i llama_sampler_greedy_i = {1093 /* .name = */ llama_sampler_greedy_name,1094 /* .accept = */ nullptr,1095 /* .apply = */ llama_sampler_greedy_apply,1096 /* .reset = */ llama_sampler_greedy_reset,1097 /* .clone = */ llama_sampler_greedy_clone,1098 /* .free = */ llama_sampler_greedy_free,1099 /* .backend_init = */ llama_sampler_greedy_backend_init,1100 /* .backend_accept = */ nullptr,1101 /* .backend_apply = */ llama_sampler_greedy_backend_apply,1102 /* .backend_set_input = */ nullptr,1103 /* .backend_reset = */ nullptr,1104 /* .copy_state = */ llama_sampler_backend_copy_state<llama_sampler_greedy>,1105};1106 1107struct llama_sampler * llama_sampler_init_greedy() {1108 return llama_sampler_init(1109 /* .iface = */ &llama_sampler_greedy_i,1110 /* .ctx = */ new llama_sampler_greedy {1111 ("greedy"),1112 }1113 );1114}1115 1116// dist1117 1118struct llama_sampler_dist : public llama_sampler_backend {1119 const uint32_t seed;1120 uint32_t seed_cur;1121 1122 std::mt19937 rng;1123 1124 // TODO: refactor + fix naming1125 // https://github.com/ggml-org/llama.cpp/pull/25532/changes#r37499067191126 // use a temporary RNG for multi-output sampling so rejected tokens do not advance rng1127 bool backend_transactional;1128 std::mt19937 rng_backend;1129 size_t n_backend_draws_generated;1130 size_t n_backend_draws_committed;1131 1132 // inputs for the current sampling graph1133 std::vector<ggml_tensor *> inp_uniforms;1134 1135 void copy_state(const llama_sampler_dist & src) {1136 // note: inp_uniforms and backend_transactional belong to the current sampling graph1137 seed_cur = src.seed_cur;1138 rng = src.rng;1139 rng_backend = src.rng_backend;1140 n_backend_draws_generated = src.n_backend_draws_generated;1141 n_backend_draws_committed = src.n_backend_draws_committed;1142 }1143};1144 1145static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {1146 auto * sctx = (llama_sampler_dist *) smpl->ctx;1147 return sctx->get_name();1148}1149 1150static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {1151 auto * ctx = (llama_sampler_dist *) smpl->ctx;1152 1153 // edge cases1154 if (cur_p->size == 0) {1155 cur_p->selected = -1;1156 return;1157 }1158 1159 cur_p->selected = 0;1160 1161 std::uniform_real_distribution<double> dist(0.0f, 1.0f);1162 1163 if (cur_p->size == 1) {1164 // keep the RNG state aligned with backend sampling, which draws once per output1165 dist(ctx->rng);1166 cur_p->data[0].p = 1.0f;1167 return;1168 }1169 1170 // max logit for numerical stability1171 float max_l = cur_p->data[0].logit;1172 if (!cur_p->sorted) {1173 for (size_t i = 1; i < cur_p->size; ++i) {1174 max_l = std::max(max_l, cur_p->data[i].logit);1175 }1176 }1177 1178 // apply softmax to obtain the probabilities1179 double sum_cum = 0.0f;1180 for (size_t i = 0; i < cur_p->size; ++i) {1181 float p = expf(cur_p->data[i].logit - max_l);1182 cur_p->data[i].p = p;1183 sum_cum += p;1184 }1185 1186#if 11187 // sample from the obtained probabilities and normalize the probs in a single pass1188 // this is ~3x faster on Mac with full gpt-oss vocab than the version below1189 //1190 const double rnd = dist(ctx->rng);1191 1192 double sum_run = 0.0f;1193 const double sum_tgt = sum_cum*rnd;1194 1195 bool found = false;1196 for (size_t i = 0; i < cur_p->size; ++i) {1197 if (!found) {1198 // accumulate probs until we reach the target sum1199 sum_run += cur_p->data[i].p;1200 if (sum_run >= sum_tgt) {