Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-18 16:55:05 +02:00
Code Issues Packages Projects Releases Wiki Activity

sampling : add support for backend sampling

Browse Source 
This commit adds support for performing sampling operations on the
backend (e.g. GPU) as part of the model computation graph.

The motivation for this feature is to enable sampling to be performed
directly on the backend as part of the computation graph being executed,
allowing for some or all of the sampling to be done on the backend.

For example, the backend sampler chain might select/sample a token
directly in which case only the sampled token needs to be transferred
from device memory to host memory.

It is also possible for the backend samplers to perform filtering of
the logits, or compute and filter the probability distribution, in
which case only the filtered logits or probabilites need to be
transferred back to system memory for further processing by CPU
samplers.

Currently the backend sampling works in a similar manner to how
pooling works, it is a function that is called by build_graph and the
sampler operations become part of the models computation graph.
This commit is contained in:
Daniel Bevenius
2025-11-17 15:19:34 +01:00
parent cb623de3fc
commit 7884b0e0ac
18 changed files with 2334 additions and 194 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/sampling.cpp
+87 -4
View File
@@ -113,23 +113,61 @@ struct common_sampler {
llama_token_data_array cur_p;
void set_logits(struct llama_context * ctx, int idx) {
const auto * logits = llama_get_logits_ith(ctx, idx);
const float * sampled_probs = llama_get_backend_sampled_probs_ith(ctx, idx);
const float * sampled_logits = llama_get_backend_sampled_logits_ith(ctx, idx);
const llama_token * sampled_ids = llama_get_backend_sampled_token_ids_ith(ctx, idx);
const llama_model * model = llama_get_model(ctx);
const llama_vocab * vocab = llama_model_get_vocab(model);
const int n_vocab = llama_vocab_n_tokens(vocab);
cur.resize(n_vocab);
// Use the member variable instead of allocating locally
cur.clear();
for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
cur[token_id] = llama_token_data{token_id, logits[token_id], 0.0f};
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]});
}
}
} 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 (llama_token i = 0; i < (int)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});
}
}
} else {
const auto * logits = llama_get_logits_ith(ctx, idx);
GGML_ASSERT(logits != nullptr);
cur.reserve(n_vocab);
for (llama_token token_id = 0; token_id < n_vocab; token_id++) {
cur.emplace_back(llama_token_data{token_id, logits[token_id], 0.0f});
}
}
cur_p = { cur.data(), cur.size(), -1, false };
}
};
static bool sampler_enabled(const struct common_params_sampling & params, enum common_sampler_type type) {
return std::find(params.samplers.begin(), params.samplers.end(), type) != params.samplers.end();
}
std::string common_params_sampling::print() const {
char result[1024];
@@ -287,6 +325,43 @@ struct common_sampler * common_sampler_init(const struct llama_model * model, co
return result;
}
struct llama_sampler * common_sampler_backend_init(const struct llama_model * model, const struct common_params_sampling & params) {
const llama_vocab * vocab = llama_model_get_vocab(model);
llama_sampler_chain_params chain_params = llama_sampler_chain_default_params();
chain_params.no_perf = params.no_perf;
struct llama_sampler * chain = llama_sampler_chain_init(chain_params);
if (!params.backend_sampling) {
return chain; // return empty chain
}
const bool enable_temp = params.temp > 0.0f && sampler_enabled(params, COMMON_SAMPLER_TYPE_TEMPERATURE);
const bool enable_top_k = params.top_k > 0 && sampler_enabled(params, COMMON_SAMPLER_TYPE_TOP_K);
const bool enable_dist = params.backend_dist;
if (!params.logit_bias.empty()) {
llama_sampler_chain_add(chain, llama_sampler_backend_init_logit_bias(
llama_vocab_n_tokens(vocab),
params.logit_bias.size(),
params.logit_bias.data()));
}
if (enable_temp) {
llama_sampler_chain_add(chain, llama_sampler_backend_init_temp(params.temp));
}
if (enable_top_k) {
llama_sampler_chain_add(chain, llama_sampler_backend_init_top_k(params.top_k));
}
if (enable_dist) {
llama_sampler_chain_add(chain, llama_sampler_backend_init_dist(params.seed));
}
return chain;
}
void common_sampler_free(struct common_sampler * gsmpl) {
if (gsmpl) {
llama_sampler_free(gsmpl->grmr);
@@ -337,6 +412,14 @@ void common_perf_print(const struct llama_context * ctx, const struct common_sam
}
llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_context * ctx, int idx, bool grammar_first) {
// Check if a backend sampler has already sampled a token in which case we
// return that token id directly.
const llama_token backend_sampled_token = llama_get_backend_sampled_token_ith(ctx, idx);
if (backend_sampled_token != LLAMA_TOKEN_NULL) {
LOG_DBG("%s: Backend sampler selected token: '%d'. Will not run any CPU samplers\n", __func__, backend_sampled_token);
return backend_sampled_token;
}
gsmpl->set_logits(ctx, idx);
auto & grmr = gsmpl->grmr;