sampling : optimize samplers by reusing bucket sort #15665
Conversation
There was a problem hiding this comment.
Logically this seems correct but I think it would be better to integrate the limit for top_p into the bucket sort function. You could keep track of the probability content per bucket and only process as many as would be needed to reach the required threshold. Though quite honestly I think the top 256 tokens are going to be enough for all practical applications so this probably doesn't matter much.
| if (post_sampling) { | ||
| const auto * cur_p = common_sampler_get_candidates(slot.smpl); | ||
| const auto * cur_p = common_sampler_get_candidates(slot.smpl, true); |
There was a problem hiding this comment.
In situations where the application requires the candidates to be sorted, using common_sampler_get_candidates(smpl, true); will perform the sorting for convenience.
ggerganov
force-pushed
the
gg/sampling-sort-opt
branch
from
37d3dd4 to
2efc7e4
Compare
src/llama-sampling.cpp
Outdated
| static void llama_token_data_array_sort_inplace(llama_token_data_array * cur_p, int k, llama_sort_data & buf) { | ||
| static const auto comp = [](const llama_token_data & a, const llama_token_data & b) { | ||
| return a.logit > b.logit; | ||
| }; | ||
|
|
||
| if (k <= 128) { | ||
| std::partial_sort(cur_p->data, cur_p->data + k, cur_p->data + cur_p->size, comp); | ||
| return; | ||
| } | ||
|
|
||
| llama_token_data_array_sort(*cur_p, k, buf); | ||
|
|
||
| std::memcpy(cur_p->data, buf.data.data(), k*sizeof(llama_token_data)); | ||
| } |
There was a problem hiding this comment.
cur_p->sorted = true could be set here, currently it is done after every call to this function.
There was a problem hiding this comment.
The last memcpy copying only the first k elements does not seem correct. If you do a partial sort, you still need to copy all the elements. If the intention is to reduce the size, then cur_p->size should be updated here.
There was a problem hiding this comment.
Also, in C++ code std::copy may be preferred to memcpy, since it is a type-safe.
There was a problem hiding this comment.
One more thing, it was not obvious to me that the k parameter is used to do a partial sort. The name of the argument is not clear enough by itself, so this should be documented. Renaming the function to "partial_sort" could also help.
ggerganov
force-pushed
the
gg/sampling-sort-opt
branch
from
d74a6ab to
6d2a38c
Compare
|
Thanks. I removed the helpers buffers and updated the code as recommended ( |
Co-authored-by: Johannes Gäßler <[email protected]>
The major change here is that the
distsampler no longer sorts. A lot of the tests intest-samplingassume that the outputs are sorted, hence the many reordering of the values there. If we don't make this change, then using a sampler chain such as the one recommended by OpenAI forgpt-oss(top-p=1 + min-p=0 + top-k=0 (i.e. disabled)) would result in very slow sorting of the full vocabulary in thedistsampler because there is nothing to cut the low-probability tokens.top-pandmin-psamplersdistsampler (255b070)common_sampler_get_candidates()can now explicitly return sorted candidates if requested viabool do_sortlibllamaAPI changesllama_sampler_init_softmax()distsampler created withllama_sampler_init_dist()will no longer sort the candidates. The old behaviour of implicitly sorting the candidates was not documented, so technically this is not a breaking change, but it's possible that user code assumed that the results will be sorted - therefore making a note