server : optimize get_token_probabilities (#24796)

Use std::partial_sort to order only the requested top-n tokens instead
of the full vocabulary

    logprobs sort: vocab=128000 n_top=0 iters=100
    full    sort:   8555.6 us/op
    partial sort:    704.3 us/op

Signed-off-by: Adrien Gallouët <angt@huggingface.co>
This commit is contained in:
Adrien Gallouët
2026-06-19 23:26:54 +02:00
committed by GitHub
parent e475fa2b5f
commit 4b48a53b6c
3 changed files with 27 additions and 14 deletions
+1 -2
View File
@@ -1824,8 +1824,7 @@ private:
});
}
} else {
// TODO: optimize this with min-p optimization
std::vector<llama_token_data> cur = get_token_probabilities(ctx_tgt, idx);
std::vector<llama_token_data> cur = get_token_probabilities(ctx_tgt, idx, n_probs_request);
const size_t max_probs = cur.size();
const size_t n_probs = std::min(max_probs, n_probs_request);