ikawrakow/ik_llama.cpp
1
0
Fork 0
mirror of https://github.com/ikawrakow/ik_llama.cpp.git synced 2026-05-01 11:51:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

sampling: refactor sorting (#1166)

Browse Source 
* sampling: refactor sorting

* Couldn't look at it without fixing it.
This commit is contained in:
Kawrakow
2026-01-19 16:48:54 +02:00
committed by GitHub
parent 98b30e5e81
commit ef5f17940c
1 changed files with 78 additions and 81 deletions
Download Patch File Download Diff File Expand all files Collapse all files

109
src/llama-sampling.cpp
View File

@@ -33,57 +33,22 @@ void llama_set_rng_seed_impl(struct llama_sampling * smpl, uint32_t seed) {
smpl->rng.seed(seed);
}
void llama_sample_softmax_impl(struct llama_sampling * smpl, llama_token_data_array * candidates) {
GGML_ASSERT(candidates->size > 0);
const int64_t t_start_sample_us = ggml_time_us();
// Sort the logits in descending order
if (!candidates->sorted) {
std::sort(candidates->data, candidates->data + candidates->size, [](const llama_token_data & a, const llama_token_data & b) {
return a.logit > b.logit;
});
candidates->sorted = true;
static void llama_sort(llama_token_data_array * candidates, int32_t k) {
if (candidates->sorted || candidates->size < 2) {
return;
}
float max_l = candidates->data[0].logit;
float cum_sum = 0.0f;
for (size_t i = 0; i < candidates->size; ++i) {
float p = expf(candidates->data[i].logit - max_l);
candidates->data[i].p = p;
cum_sum += p;
}
for (size_t i = 0; i < candidates->size; ++i) {
candidates->data[i].p /= cum_sum;
}
if (smpl) {
smpl->t_sample_us += ggml_time_us() - t_start_sample_us;
}
}
void llama_sample_top_k_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, int32_t k, size_t min_keep) {
// TODO: move bucket sort to separate function so that top_p/tail_free/typical/softmax first is equally fast
// if (k >= (int32_t)candidates->size) {
// return;
// }
const int64_t t_start_sample_us = ggml_time_us();
if (k <= 0) {
if (k < 0) {
k = candidates->size;
}
k = std::max(k, (int) min_keep);
k = std::min(k, (int) candidates->size);
// Sort scores in descending order
if (!candidates->sorted) {
auto comp = [](const llama_token_data & a, const llama_token_data & b) {
return a.logit > b.logit;
};
if (k <= 128) {
if (k <= 1024) { //128) {
if (k == int(candidates->size)) {
std::sort(candidates->data, candidates->data + candidates->size, comp);
} else {
std::partial_sort(candidates->data, candidates->data + k, candidates->data + candidates->size, comp);
}
} else {
constexpr int nbuckets = 128;
constexpr float bucket_low = -10.0f;
@@ -136,6 +101,44 @@ void llama_sample_top_k_impl(struct llama_sampling * smpl, llama_token_data_arra
}
candidates->sorted = true;
}
void llama_sample_softmax_impl(struct llama_sampling * smpl, llama_token_data_array * candidates) {
GGML_ASSERT(candidates->size > 0);
const int64_t t_start_sample_us = ggml_time_us();
// Sort the logits in descending order if necessary
llama_sort(candidates, -1);
float max_l = candidates->data[0].logit;
float cum_sum = 0.0f;
for (size_t i = 0; i < candidates->size; ++i) {
float p = expf(candidates->data[i].logit - max_l);
candidates->data[i].p = p;
cum_sum += p;
}
for (size_t i = 0; i < candidates->size; ++i) {
candidates->data[i].p /= cum_sum;
}
if (smpl) {
smpl->t_sample_us += ggml_time_us() - t_start_sample_us;
}
}
void llama_sample_top_k_impl(struct llama_sampling * smpl, llama_token_data_array * candidates, int32_t k, size_t min_keep) {
const int64_t t_start_sample_us = ggml_time_us();
if (k <= 0) {
k = candidates->size;
}
k = std::max(k, (int) min_keep);
k = std::min(k, (int) candidates->size);
llama_sort(candidates, k);
candidates->size = k;
if (smpl) {
@@ -210,13 +213,8 @@ void llama_sample_min_p_impl(struct llama_sampling * smpl, llama_token_data_arra
// if the candidates are sorted or the unsorted implementation failed, use this implementation
if (!min_p_applied) {
// Sort the logits in descending order
if (!candidates->sorted) {
std::sort(candidates->data, candidates->data + candidates->size, [](const llama_token_data & a, const llama_token_data & b) {
return a.logit > b.logit;
});
candidates->sorted = true;
}
// Sort the logits in descending order if needed
llama_sort(candidates, -1);
const float min_logit = candidates->data[0].logit + logf(p); // min logit for p_i >= p * p_max
size_t i = 1; // first token always matches
@@ -313,10 +311,9 @@ void llama_sample_typical_impl(struct llama_sampling * smpl, llama_token_data_ar
}
// Compute the absolute difference between negative log probability and entropy for each candidate
std::vector<float> shifted_scores;
std::vector<float> shifted_scores(candidates->size);
for (size_t i = 0; i < candidates->size; ++i) {
float shifted_score = fabsf(-logf(candidates->data[i].p) - entropy);
shifted_scores.push_back(shifted_score);
shifted_scores[i] = fabsf(-logf(candidates->data[i].p) - entropy);
}
// Sort tokens based on the shifted_scores and their corresponding indices
@@ -343,10 +340,10 @@ void llama_sample_typical_impl(struct llama_sampling * smpl, llama_token_data_ar
}
// Resize the output vector to keep only the locally typical tokens
std::vector<llama_token_data> new_candidates;
std::vector<llama_token_data> new_candidates(last_idx);
for (size_t i = 0; i < last_idx; ++i) {
size_t idx = indices[i];
new_candidates.push_back(candidates->data[idx]);
new_candidates[i] = candidates->data[idx];
}
// Replace the data in candidates with the new_candidates data