nikola/llama_cpp
1
0
Fork 0
Code Issues Pull Requests Actions 23 Packages Projects Releases Wiki Activity

server: fix memory reservations in populate_token_probs (#18787)

Browse Source
This commit is contained in:
Lennart Austenfeld
2026-01-19 19:13:31 +01:00
committed by GitHub
parent 365a3e8c31
commit 18361c579c
1 changed files with 8 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/server/server-context.cpp
+8 -5
View File
@@ -1326,11 +1326,12 @@ private:
} }
void populate_token_probs(const server_slot & slot, completion_token_output & result, bool post_sampling, bool special, int idx) const { void populate_token_probs(const server_slot & slot, completion_token_output & result, bool post_sampling, bool special, int idx) const {
const size_t n_probs = slot.task->params.sampling.n_probs; const size_t n_probs_request = slot.task->params.sampling.n_probs;
if (post_sampling) { if (post_sampling) {
const auto * cur_p = common_sampler_get_candidates(slot.smpl.get(), true); const auto * cur_p = common_sampler_get_candidates(slot.smpl.get(), true);
const size_t max_probs = cur_p->size; const size_t max_probs = cur_p->size;
const size_t n_probs = std::min(max_probs, n_probs_request);
// set probability for sampled token // set probability for sampled token
for (size_t i = 0; i < max_probs; i++) { for (size_t i = 0; i < max_probs; i++) {
@@ -1341,8 +1342,8 @@ private:
} }
// set probability for top n_probs tokens // set probability for top n_probs tokens
result.probs.reserve(max_probs); result.probs.reserve(n_probs);
for (size_t i = 0; i < std::min(max_probs, n_probs); i++) { for (size_t i = 0; i < n_probs; i++) {
result.probs.push_back({ result.probs.push_back({
cur_p->data[i].id, cur_p->data[i].id,
common_token_to_piece(ctx, cur_p->data[i].id, special), common_token_to_piece(ctx, cur_p->data[i].id, special),
@@ -1352,9 +1353,11 @@ private:
} else { } else {
// TODO: optimize this with min-p optimization // TODO: optimize this with min-p optimization
std::vector<llama_token_data> cur = get_token_probabilities(ctx, idx); std::vector<llama_token_data> cur = get_token_probabilities(ctx, idx);
const size_t max_probs = cur.size();
const size_t n_probs = std::min(max_probs, n_probs_request);
// set probability for sampled token // set probability for sampled token
for (size_t i = 0; i < cur.size(); i++) { for (size_t i = 0; i < max_probs; i++) {
// set probability for sampled token // set probability for sampled token
if (cur[i].id == result.tok) { if (cur[i].id == result.tok) {
result.prob = cur[i].p; result.prob = cur[i].p;
@@ -1364,7 +1367,7 @@ private:
// set probability for top n_probs tokens // set probability for top n_probs tokens
result.probs.reserve(n_probs); result.probs.reserve(n_probs);
for (size_t i = 0; i < std::min(cur.size(), n_probs); i++) { for (size_t i = 0; i < n_probs; i++) {
result.probs.push_back({ result.probs.push_back({
cur[i].id, cur[i].id,
common_token_to_piece(ctx, cur[i].id, special), common_token_to_piece(ctx, cur[i].id, special),