vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Bucket sort: another minor improvement

Browse source
This commit is contained in:
Iwan Kawrakow 2024-01-24 15:03:07 +02:00
parent e7c1f64c50
commit 26dde91aa6
1 changed files with 12 additions and 10 deletions
Download patch file Download diff file Expand all files Collapse all files

16
llama.cpp
View file

@ -8001,7 +8001,7 @@ void llama_sample_top_k(struct llama_context * ctx, llama_token_data_array * can
auto comp = [](const llama_token_data & a, const llama_token_data & b) { auto comp = [](const llama_token_data & a, const llama_token_data & b) {
return a.logit > b.logit; return a.logit > b.logit;
}; };
if (k < 200) { if (k <= 128) {
std::partial_sort(candidates->data, candidates->data + k, candidates->data + candidates->size, comp); std::partial_sort(candidates->data, candidates->data + k, candidates->data + candidates->size, comp);
} else { } else {
constexpr int nbuckets = 128; constexpr int nbuckets = 128;
@ -8010,11 +8010,8 @@ void llama_sample_top_k(struct llama_context * ctx, llama_token_data_array * can
constexpr float bucket_scale = nbuckets/(bucket_high - bucket_low); constexpr float bucket_scale = nbuckets/(bucket_high - bucket_low);
constexpr float bucker_inter = -bucket_low * bucket_scale; constexpr float bucker_inter = -bucket_low * bucket_scale;
std::vector<llama_token_data> tmp_tokens(candidates->size);
std::vector<int> bucket_idx(candidates->size); std::vector<int> bucket_idx(candidates->size);
std::vector<int> histo(nbuckets, 0); std::vector<int> histo(nbuckets, 0);
std::vector<llama_token_data*> bucket_ptrs;
bucket_ptrs.reserve(nbuckets);
for (int i = 0; i < (int)candidates->size; ++i) { for (int i = 0; i < (int)candidates->size; ++i) {
const float val = candidates->data[i].logit; const float val = candidates->data[i].logit;
@ -8023,14 +8020,19 @@ void llama_sample_top_k(struct llama_context * ctx, llama_token_data_array * can
bucket_idx[i] = ib; bucket_idx[i] = ib;
++histo[ib]; ++histo[ib];
} }
auto ptr = tmp_tokens.data();
int nhave = 0; int nhave = 0;
int ib = nbuckets - 1; int ib = nbuckets - 1;
for ( ; ib >= 0; --ib) { for ( ; ib >= 0; --ib) {
nhave += histo[ib]; nhave += histo[ib];
bucket_ptrs.push_back(ptr);
if (nhave >= k) break; if (nhave >= k) break;
ptr += histo[ib]; }
std::vector<llama_token_data> tmp_tokens(nhave);
auto ptr = tmp_tokens.data();
std::vector<llama_token_data*> bucket_ptrs;
bucket_ptrs.reserve(nbuckets - ib);
for (int j = nbuckets - 1; j >= ib; --j) {
bucket_ptrs.push_back(ptr);
ptr += histo[j];
} }
for (int i = 0; i < (int)candidates->size; ++i) { for (int i = 0; i < (int)candidates->size; ++i) {
int j = bucket_idx[i]; int j = bucket_idx[i];