sampling : always expose sampled_ids

This commit precomputes and caches the full-vocab token id list in
llama_context's constructor, so llama_get_backend_sampled_token_ids_ith
always returns a valid pointer.

The motivation for this is that this enables both common/sampling.cpp
and src/llama-sampling.cpp can simplify their logic.

Not all backends samplers that process logits need to set the
sampled_tokens_id as they may not change the order of the logits, for
example the temperature sampler only scales the logits but does not
change their order. Simliar the logit bias sampler only adds bias to
specific token ids but does not change the order of the logits. In
these cases there will not be a device to host copy of the sampled
token ids, and this is the use case where having this precomputed
list is useful.
This commit is contained in:
Daniel Bevenius
2025-11-18 14:54:49 +01:00
parent 4b52e59903
commit 82957a90f2
4 changed files with 26 additions and 40 deletions
+4 -18
View File
@@ -128,28 +128,14 @@ struct common_sampler {
if (sampled_probs) {
const uint32_t sampled_probs_count = llama_get_backend_sampled_probs_count_ith(ctx, idx);
cur.reserve(sampled_probs_count);
// The backend sampler has filtered the probabilities so we need to use the sampled ids.
if (sampled_ids != nullptr) {
for (uint32_t i = 0; i < sampled_probs_count; ++i) {
cur.emplace_back(llama_token_data{sampled_ids[i], 0.0f, sampled_probs[i]});
}
} else {
for (llama_token token_id = 0; token_id < (int) sampled_probs_count; token_id++) {
cur.emplace_back(llama_token_data{token_id, 0.0f, sampled_probs[token_id]});
}
for (uint32_t i = 0; i < sampled_probs_count; ++i) {
cur.emplace_back(llama_token_data{sampled_ids[i], 0.0f, sampled_probs[i]});
}
} else if (sampled_logits) {
const uint32_t sampled_logits_count = llama_get_backend_sampled_logits_count_ith(ctx, idx);
cur.reserve(sampled_logits_count);
// The backend sampler has filtered the logits so we need to use the sampled ids.
if (sampled_ids != nullptr) {
for (uint32_t i = 0; i < sampled_logits_count; i++) {
cur.emplace_back(llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f});
}
} else {
for (llama_token token_id = 0; token_id < (int) sampled_logits_count; token_id++) {
cur.emplace_back(llama_token_data{token_id, sampled_logits[token_id], 0.0f});
}
for (uint32_t i = 0; i < sampled_logits_count; i++) {
cur.emplace_back(llama_token_data{sampled_ids[i], sampled_logits[i], 0.0f});
}
} else {
const auto * logits = llama_get_logits_ith(ctx, idx);