mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-08 22:09:02 +02:00
Enable backend sampling with token speculation
This commit is contained in:
@@ -115,6 +115,7 @@ static llama_sampler_i llama_sampler_llg_i = {
|
||||
/* .backend_init = */ NULL,
|
||||
/* .backend_accept = */ NULL,
|
||||
/* .backend_apply = */ NULL,
|
||||
/* .backend_reset = */ NULL,
|
||||
/* .backend_set_input = */ NULL,
|
||||
};
|
||||
|
||||
|
||||
@@ -216,6 +216,7 @@ static struct llama_sampler_i common_reasoning_budget_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
|
||||
@@ -2377,6 +2377,11 @@ common_speculative_init_result_ptr common_speculative_init_from_params(common_pa
|
||||
return std::make_unique<common_speculative_init_result>(params, model_tgt, ctx_tgt);
|
||||
}
|
||||
|
||||
int32_t common_speculative_n_outputs_max(int32_t n_batch, int32_t n_parallel, int32_t n_draft) {
|
||||
const int64_t n_outputs = (int64_t) n_parallel * (1 + (int64_t) std::max(0, n_draft));
|
||||
return std::min<int64_t>(n_batch, n_outputs);
|
||||
}
|
||||
|
||||
// initialization of the speculative decoding system
|
||||
//
|
||||
common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {
|
||||
|
||||
@@ -25,6 +25,9 @@ int32_t common_speculative_n_max(const common_params_speculative * spec);
|
||||
|
||||
common_params common_base_params_to_speculative(const common_params & params);
|
||||
|
||||
// return the max number of outputs needed for speculative decoding
|
||||
int32_t common_speculative_n_outputs_max(int32_t n_batch, int32_t n_parallel, int32_t n_draft);
|
||||
|
||||
common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq);
|
||||
|
||||
void common_speculative_free(common_speculative * spec);
|
||||
|
||||
@@ -3,9 +3,11 @@
|
||||
#include "common.h"
|
||||
#include "ngram-cache.h"
|
||||
#include "sampling.h"
|
||||
#include "speculative.h"
|
||||
#include "log.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <clocale>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
@@ -27,6 +29,8 @@ int main(int argc, char ** argv){
|
||||
// max. number of additional tokens to draft if match is found
|
||||
const int n_draft = params.speculative.draft.n_max;
|
||||
|
||||
params.n_outputs_max = common_speculative_n_outputs_max(params.n_batch, params.n_parallel, n_draft);
|
||||
|
||||
// init llama.cpp
|
||||
llama_backend_init();
|
||||
llama_numa_init(params.numa);
|
||||
|
||||
@@ -5,6 +5,7 @@
|
||||
#include "log.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <clocale>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
@@ -29,6 +30,9 @@ int main(int argc, char ** argv) {
|
||||
return 1;
|
||||
}
|
||||
|
||||
params.n_outputs_max = common_speculative_n_outputs_max(
|
||||
params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative));
|
||||
|
||||
// init llama.cpp
|
||||
llama_backend_init();
|
||||
llama_numa_init(params.numa);
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#include "arg.h"
|
||||
#include "common.h"
|
||||
#include "sampling.h"
|
||||
#include "speculative.h"
|
||||
#include "log.h"
|
||||
#include "llama.h"
|
||||
|
||||
@@ -57,6 +58,9 @@ int main(int argc, char ** argv) {
|
||||
// max number of parallel drafting sequences (i.e. tree branches)
|
||||
const int n_seq_dft = params.n_parallel;
|
||||
|
||||
params.n_outputs_max = common_speculative_n_outputs_max(
|
||||
params.n_batch, params.n_parallel, params.speculative.draft.n_max);
|
||||
|
||||
// probability threshold for splitting a draft branch (only for n_seq_dft > 1)
|
||||
const float p_draft_split = params.speculative.draft.p_split;
|
||||
|
||||
|
||||
+8
-2
@@ -1270,9 +1270,12 @@ extern "C" {
|
||||
// [EXPERIMENTAL]
|
||||
// backend sampling interface:
|
||||
|
||||
// return true if the backend supports all ops needed by the sampler
|
||||
// return true if the backend supports all ops needed by the sampler and the requested output mode
|
||||
// note: call once per sampler
|
||||
bool (*backend_init)(struct llama_sampler * smpl, ggml_backend_buffer_type_t buft);
|
||||
bool (*backend_init)(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output);
|
||||
|
||||
// call after .backend_apply()
|
||||
void (*backend_accept)(
|
||||
@@ -1288,6 +1291,9 @@ extern "C" {
|
||||
struct ggml_cgraph * gf,
|
||||
struct llama_sampler_data * data);
|
||||
|
||||
// called before rebuilding a sampling graph to clear graph-owned tensor references
|
||||
void (*backend_reset)(struct llama_sampler * smpl);
|
||||
|
||||
// called before graph execution to set inputs for the current ubatch
|
||||
void (*backend_set_input)(struct llama_sampler * smpl);
|
||||
};
|
||||
|
||||
+77
-112
@@ -10,6 +10,7 @@
|
||||
#include "llama-mmap.h"
|
||||
#include "llama-model.h"
|
||||
#include "llama-ext.h"
|
||||
#include "llama-sampler.h"
|
||||
#include "llama.h"
|
||||
|
||||
#include <cinttypes>
|
||||
@@ -99,6 +100,7 @@ llama_context::llama_context(
|
||||
if (cparams.n_seq_max > LLAMA_MAX_SEQ) {
|
||||
throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ));
|
||||
}
|
||||
cparams.sampler_backend_require_multi_output = params.n_outputs_max > cparams.n_seq_max;
|
||||
|
||||
cparams.n_rs_seq = params.n_rs_seq;
|
||||
if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
|
||||
@@ -1231,7 +1233,7 @@ bool llama_context::set_sampler(llama_seq_id seq_id, llama_sampler * sampler) {
|
||||
if (sampler && can_offload) {
|
||||
auto * buft = ggml_backend_dev_buffer_type(model.dev_output());
|
||||
|
||||
sampler->iface->backend_init(sampler, buft);
|
||||
sampler->iface->backend_init(sampler, buft, cparams.sampler_backend_require_multi_output);
|
||||
|
||||
sampling.samplers[seq_id] = sampler;
|
||||
|
||||
@@ -1576,108 +1578,38 @@ int llama_context::encode(const llama_batch & batch_inp) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
static std::map<llama_seq_id, uint32_t> build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) {
|
||||
std::map<llama_seq_id, uint32_t> seq_to_row;
|
||||
// how many output tokens we have seen so far for this ubatch.
|
||||
uint32_t local = 0;
|
||||
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
|
||||
// skip tokens that are not output.
|
||||
if (!ubatch.output[i]) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const llama_seq_id seq_id = ubatch.seq_id[i][0];
|
||||
// row_offset is the number of output tokens before this ubatch.
|
||||
seq_to_row[seq_id] = row_offset + local;
|
||||
++local;
|
||||
}
|
||||
return seq_to_row;
|
||||
}
|
||||
|
||||
static void copy_tensor_async_ints(
|
||||
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
|
||||
const buffer_view<llama_token> & sampled,
|
||||
const std::map<llama_seq_id, uint32_t> & seq_to_row,
|
||||
ggml_backend_sched_t sched) {
|
||||
if (!sampled.has_data()) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (const auto & [seq_id, tensor] : tensor_map) {
|
||||
auto it = seq_to_row.find(seq_id);
|
||||
if (it == seq_to_row.end()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const uint32_t row = it->second;
|
||||
GGML_ASSERT(row < sampled.size);
|
||||
|
||||
GGML_ASSERT(ggml_is_contiguous(tensor) && "sampled tokens tensor must be contiguous for async copy");
|
||||
|
||||
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
|
||||
ggml_backend_tensor_get_async(backend, tensor, sampled.data + row, 0, sizeof(sampled.data[row]));
|
||||
}
|
||||
}
|
||||
|
||||
static void copy_tensor_async_floats(
|
||||
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
|
||||
const buffer_view<float> & dst,
|
||||
template<typename T>
|
||||
static void copy_tensor_async_rows(
|
||||
const std::vector<ggml_tensor *> & tensors,
|
||||
const buffer_view<T> & dst,
|
||||
size_t stride,
|
||||
std::vector<uint32_t> & counts,
|
||||
const std::map<llama_seq_id, uint32_t> & seq_to_row,
|
||||
ggml_backend_sched_t sched) {
|
||||
uint32_t row_offset,
|
||||
ggml_backend_sched_t sched,
|
||||
std::vector<uint32_t> * counts = nullptr) {
|
||||
if (!dst.has_data()) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (const auto & [seq_id, tensor] : tensor_map) {
|
||||
auto it = seq_to_row.find(seq_id);
|
||||
if (it == seq_to_row.end()) {
|
||||
for (size_t i = 0; i < tensors.size(); ++i) {
|
||||
auto * tensor = tensors[i];
|
||||
if (tensor == nullptr) {
|
||||
continue;
|
||||
}
|
||||
|
||||
const uint32_t row = it->second;
|
||||
GGML_ASSERT(row < counts.size());
|
||||
|
||||
GGML_ASSERT(ggml_is_contiguous(tensor) && "logits/probs tensor must be contiguous for async copy");
|
||||
const uint32_t row = row_offset + i;
|
||||
const size_t n_elements = ggml_nelements(tensor);
|
||||
GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy");
|
||||
GGML_ASSERT(n_elements <= stride);
|
||||
GGML_ASSERT((size_t) row * stride + n_elements <= dst.size);
|
||||
|
||||
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
|
||||
float * row_ptr = dst.data + (size_t) row * stride;
|
||||
T * row_ptr = dst.data + (size_t) row * stride;
|
||||
ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
|
||||
|
||||
// Update the actual number of logits/probabilities that were written for this row.
|
||||
counts[row] = ggml_nelements(tensor);
|
||||
}
|
||||
}
|
||||
|
||||
static void copy_tensor_async_candidates(
|
||||
const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
|
||||
const buffer_view<llama_token> & dst,
|
||||
size_t stride,
|
||||
std::vector<uint32_t> & counts,
|
||||
const std::map<llama_seq_id, uint32_t> & seq_to_row,
|
||||
ggml_backend_sched_t sched) {
|
||||
if (!dst.has_data()) {
|
||||
return;
|
||||
}
|
||||
|
||||
for (const auto & [seq_id, tensor] : tensor_map) {
|
||||
auto it = seq_to_row.find(seq_id);
|
||||
if (it == seq_to_row.end()) {
|
||||
continue;
|
||||
if (counts) {
|
||||
GGML_ASSERT(row < counts->size());
|
||||
(*counts)[row] = n_elements;
|
||||
}
|
||||
|
||||
const uint32_t row = it->second;
|
||||
GGML_ASSERT(row < counts.size());
|
||||
|
||||
GGML_ASSERT(ggml_is_contiguous(tensor) && "candidates tensor must be contiguous for async copy");
|
||||
|
||||
ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
|
||||
llama_token * row_ptr = dst.data + (size_t) row * stride;
|
||||
ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
|
||||
|
||||
// Update the actual number of candidates that were written.
|
||||
counts[row] = ggml_nelements(tensor);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1726,12 +1658,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
|
||||
|
||||
const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;
|
||||
|
||||
// TODO: avoid this workaround in the future
|
||||
if (has_samplers && batch_inp.logits) {
|
||||
// embedding contexts output every token even when batch.logits is not set
|
||||
if (has_samplers && (output_all || batch_inp.logits)) {
|
||||
std::vector<int32_t> seq_output_count(n_seq_max, 0);
|
||||
|
||||
for (int32_t i = 0; i < batch_inp.n_tokens; ++i) {
|
||||
if (batch_inp.logits[i] == 0) {
|
||||
if (!output_all && batch_inp.logits[i] == 0) {
|
||||
continue;
|
||||
}
|
||||
|
||||
@@ -1741,7 +1673,10 @@ int llama_context::decode(const llama_batch & batch_inp) {
|
||||
const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0;
|
||||
|
||||
seq_output_count[seq_id]++;
|
||||
if (seq_output_count[seq_id] > 1) {
|
||||
auto sampler = sampling.samplers.find(seq_id);
|
||||
if (seq_output_count[seq_id] > 1 &&
|
||||
sampler != sampling.samplers.end() &&
|
||||
!cparams.sampler_backend_require_multi_output) {
|
||||
LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n",
|
||||
__func__, seq_id, seq_output_count[seq_id]);
|
||||
return -1;
|
||||
@@ -2009,17 +1944,14 @@ int llama_context::decode(const llama_batch & batch_inp) {
|
||||
}
|
||||
}
|
||||
|
||||
// Copy backend sampling output if this ubatch produced any sampling tensors.
|
||||
if (has_samplers && (!res->t_sampled.empty() || !res->t_sampled_probs.empty() || !res->t_sampled_logits.empty())) {
|
||||
const auto seq_to_output_row = build_seq_to_output_row(ubatch, n_outputs_prev);
|
||||
if (has_samplers) {
|
||||
const auto stride = n_vocab;
|
||||
|
||||
// async copy the sampling data from the backend to the host
|
||||
copy_tensor_async_ints(res->t_sampled, sampling.sampled, seq_to_output_row, sched.get());
|
||||
|
||||
copy_tensor_async_floats (res->t_sampled_logits, sampling.logits, stride, sampling.logits_count, seq_to_output_row, sched.get());
|
||||
copy_tensor_async_floats (res->t_sampled_probs, sampling.probs, stride, sampling.probs_count, seq_to_output_row, sched.get());
|
||||
copy_tensor_async_candidates(res->t_candidates, sampling.candidates, stride, sampling.candidates_count, seq_to_output_row, sched.get());
|
||||
copy_tensor_async_rows(res->t_sampled, sampling.sampled, 1, n_outputs_prev, sched.get());
|
||||
copy_tensor_async_rows(res->t_sampled_logits, sampling.logits, stride, n_outputs_prev, sched.get(), &sampling.logits_count);
|
||||
copy_tensor_async_rows(res->t_sampled_probs, sampling.probs, stride, n_outputs_prev, sched.get(), &sampling.probs_count);
|
||||
copy_tensor_async_rows(res->t_candidates, sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count);
|
||||
}
|
||||
|
||||
n_outputs_prev += n_outputs;
|
||||
@@ -2349,6 +2281,7 @@ void llama_context::output_reorder() {
|
||||
//
|
||||
|
||||
uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
|
||||
uint32_t res;
|
||||
if (model.arch == LLM_ARCH_QWEN3NEXT ||
|
||||
model.arch == LLM_ARCH_KIMI_LINEAR ||
|
||||
model.arch == LLM_ARCH_QWEN35 ||
|
||||
@@ -2357,11 +2290,27 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
|
||||
(model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
|
||||
model.arch == LLM_ARCH_NANBEIGE ||
|
||||
model.arch == LLM_ARCH_MINIMAX_M3) {
|
||||
return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
|
||||
res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
|
||||
} else {
|
||||
res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
|
||||
for (const auto & lora : model.loras) {
|
||||
res += lora->get_n_nodes();
|
||||
}
|
||||
}
|
||||
uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
|
||||
for (const auto & lora : model.loras) {
|
||||
res += lora->get_n_nodes();
|
||||
|
||||
uint32_t n_sampling_nodes = 0;
|
||||
uint32_t n_sampling_nodes_max = 0;
|
||||
for (const auto & [seq_id, sampler] : sampling.samplers) {
|
||||
const uint32_t n_nodes = llama_sampler_backend_n_nodes(sampler);
|
||||
n_sampling_nodes += n_nodes;
|
||||
if (cparams.sampler_backend_require_multi_output) {
|
||||
n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes);
|
||||
}
|
||||
}
|
||||
|
||||
res += n_sampling_nodes;
|
||||
if (cparams.n_outputs_max > 1) {
|
||||
res += (cparams.n_outputs_max - 1) * n_sampling_nodes_max;
|
||||
}
|
||||
return res;
|
||||
}
|
||||
@@ -2394,13 +2343,29 @@ ggml_cgraph * llama_context::graph_reserve(
|
||||
llama_batch_allocr balloc(model.hparams.n_pos_per_embd());
|
||||
llama_ubatch ubatch = balloc.ubatch_reserve(n_tokens/n_seqs, n_seqs);
|
||||
|
||||
// set one output token per sequence in order to activate all backend samplers
|
||||
// spread the outputs across all sequences to reserve the largest sampling graph
|
||||
std::vector<llama_seq_id> seq_ids(n_seqs);
|
||||
for (uint32_t i = 0; i < n_seqs; ++i) {
|
||||
seq_ids[i] = i;
|
||||
ubatch.n_seq_id[i] = 1;
|
||||
ubatch.seq_id[i] = &seq_ids[i];
|
||||
ubatch.output[i] = true;
|
||||
for (uint32_t s = 0; s < n_seqs; ++s) {
|
||||
seq_ids[s] = s;
|
||||
for (uint32_t t = 0; t < ubatch.n_seq_tokens; ++t) {
|
||||
const uint32_t i = s * ubatch.n_seq_tokens + t;
|
||||
ubatch.n_seq_id[i] = 1;
|
||||
ubatch.seq_id[i] = &seq_ids[s];
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t n_outputs_set = 0;
|
||||
for (uint32_t t = 0; t < ubatch.n_seq_tokens && n_outputs_set < n_outputs; ++t) {
|
||||
for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) {
|
||||
const auto sampler = sampling.samplers.find(s);
|
||||
if (t > 0 && (sampler == sampling.samplers.end() ||
|
||||
!cparams.sampler_backend_require_multi_output)) {
|
||||
continue;
|
||||
}
|
||||
|
||||
ubatch.output[s * ubatch.n_seq_tokens + t] = true;
|
||||
++n_outputs_set;
|
||||
}
|
||||
}
|
||||
|
||||
auto * res = gf_res_reserve.get();
|
||||
|
||||
@@ -52,6 +52,7 @@ struct llama_cparams {
|
||||
bool op_offload;
|
||||
bool kv_unified;
|
||||
bool pipeline_parallel;
|
||||
bool sampler_backend_require_multi_output;
|
||||
|
||||
std::vector<bool> embeddings_layer_inp; // [n_layer()] extract input embeddings for layer
|
||||
|
||||
|
||||
+95
-69
@@ -4,6 +4,7 @@
|
||||
#include "llama-model.h"
|
||||
#include "llama-batch.h"
|
||||
#include "llama-cparams.h"
|
||||
#include "llama-sampler.h"
|
||||
|
||||
#include "llama-kv-cache.h"
|
||||
#include "llama-kv-cache-iswa.h"
|
||||
@@ -1353,24 +1354,24 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
|
||||
}
|
||||
}
|
||||
}
|
||||
for (auto & [seq_id, t] : t_sampled) {
|
||||
if (t != nullptr) {
|
||||
ggml_set_output(t);
|
||||
for (auto * tensor : t_sampled) {
|
||||
if (tensor != nullptr) {
|
||||
ggml_set_output(tensor);
|
||||
}
|
||||
}
|
||||
for (auto & [seq_id, t] : t_sampled_probs) {
|
||||
if (t != nullptr) {
|
||||
ggml_set_output(t);
|
||||
for (auto * tensor : t_sampled_probs) {
|
||||
if (tensor != nullptr) {
|
||||
ggml_set_output(tensor);
|
||||
}
|
||||
}
|
||||
for (auto & [seq_id, t] : t_sampled_logits) {
|
||||
if (t != nullptr) {
|
||||
ggml_set_output(t);
|
||||
for (auto * tensor : t_sampled_logits) {
|
||||
if (tensor != nullptr) {
|
||||
ggml_set_output(tensor);
|
||||
}
|
||||
}
|
||||
for (auto & [seq_id, t] : t_candidates) {
|
||||
if (t != nullptr) {
|
||||
ggml_set_output(t);
|
||||
for (auto * tensor : t_candidates) {
|
||||
if (tensor != nullptr) {
|
||||
ggml_set_output(tensor);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3649,77 +3650,102 @@ void llm_graph_context::build_sampling() const {
|
||||
auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);
|
||||
res->add_input(std::move(inp_sampling));
|
||||
|
||||
std::map<llama_seq_id, int32_t> seq_to_logit_row;
|
||||
int32_t logit_row_idx = 0;
|
||||
|
||||
for (uint32_t i = 0; i < ubatch.n_tokens; i++) {
|
||||
std::map<llama_seq_id, std::vector<uint32_t>> sampling_rows;
|
||||
uint32_t n_rows = 0;
|
||||
for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
|
||||
if (ubatch.output[i]) {
|
||||
llama_seq_id seq_id = ubatch.seq_id[i][0];
|
||||
seq_to_logit_row[seq_id] = logit_row_idx;
|
||||
logit_row_idx++;
|
||||
sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++);
|
||||
}
|
||||
}
|
||||
|
||||
res->t_sampled.resize(n_rows, nullptr);
|
||||
res->t_sampled_probs.resize(n_rows, nullptr);
|
||||
res->t_sampled_logits.resize(n_rows, nullptr);
|
||||
res->t_candidates.resize(n_rows, nullptr);
|
||||
|
||||
// res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1)
|
||||
GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor");
|
||||
|
||||
// add a dummy row of logits
|
||||
// this trick makes the graph static, regardless of which samplers are activated
|
||||
// this is important in order to minimize graph reallocations
|
||||
// add a dummy row to keep the single-output graph static regardless of active samplers
|
||||
// multi-output graphs can still vary with the number of output rows
|
||||
ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);
|
||||
|
||||
for (const auto & [seq_id, sampler] : samplers) {
|
||||
const auto it = seq_to_logit_row.find(seq_id);
|
||||
|
||||
// inactive samplers always work on the first row
|
||||
const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0;
|
||||
const int i_out = it != seq_to_logit_row.end() ? 1 : 0;
|
||||
|
||||
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);
|
||||
ggml_format_name(logits_seq, "logits_seq_%d", seq_id);
|
||||
|
||||
struct llama_sampler_data data = {
|
||||
/*.logits =*/ logits_seq,
|
||||
/*.probs =*/ nullptr,
|
||||
/*.sampled =*/ nullptr,
|
||||
/*.candidates =*/ nullptr,
|
||||
};
|
||||
|
||||
assert(sampler->iface->backend_apply);
|
||||
sampler->iface->backend_apply(sampler, ctx0, gf, &data);
|
||||
|
||||
if (data.sampled != nullptr) {
|
||||
res->t_sampled[seq_id] = data.sampled;
|
||||
outs[1] = data.sampled;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
|
||||
if (data.probs != nullptr) {
|
||||
res->t_sampled_probs[seq_id] = data.probs;
|
||||
outs[1] = data.probs;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
|
||||
if (data.logits != nullptr) {
|
||||
res->t_sampled_logits[seq_id] = data.logits;
|
||||
outs[1] = data.logits;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
|
||||
if (data.candidates != nullptr) {
|
||||
res->t_candidates[seq_id] = data.candidates;
|
||||
outs[1] = data.candidates;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
for (const auto & entry : samplers) {
|
||||
if (entry.second->iface->backend_reset) {
|
||||
entry.second->iface->backend_reset(entry.second);
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: Call llama_sampler_accept_ggml after all samplers have been applied.
|
||||
static const std::vector<uint32_t> dummy_row = { 0 };
|
||||
|
||||
for (const auto & [seq_id, sampler] : samplers) {
|
||||
const auto it = sampling_rows.find(seq_id);
|
||||
|
||||
// inactive samplers always work on the first row
|
||||
const bool active = it != sampling_rows.end();
|
||||
const auto & rows = active ? it->second : dummy_row;
|
||||
const int i_out = active ? 1 : 0;
|
||||
|
||||
for (uint32_t i = 0; i < rows.size(); ++i) {
|
||||
ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]);
|
||||
ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i);
|
||||
|
||||
struct llama_sampler_data data = {
|
||||
/*.logits =*/ logits_seq,
|
||||
/*.probs =*/ nullptr,
|
||||
/*.sampled =*/ nullptr,
|
||||
/*.candidates =*/ nullptr,
|
||||
};
|
||||
|
||||
assert(sampler->iface->backend_apply);
|
||||
sampler->iface->backend_apply(sampler, ctx0, gf, &data);
|
||||
|
||||
if (data.sampled != nullptr) {
|
||||
if (active) {
|
||||
res->t_sampled[rows[i]] = data.sampled;
|
||||
}
|
||||
outs[1] = data.sampled;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
|
||||
if (data.probs != nullptr) {
|
||||
if (active) {
|
||||
res->t_sampled_probs[rows[i]] = data.probs;
|
||||
}
|
||||
outs[1] = data.probs;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
|
||||
if (data.logits != nullptr) {
|
||||
if (active) {
|
||||
res->t_sampled_logits[rows[i]] = data.logits;
|
||||
}
|
||||
outs[1] = data.logits;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
|
||||
if (data.candidates != nullptr) {
|
||||
if (active) {
|
||||
res->t_candidates[rows[i]] = data.candidates;
|
||||
}
|
||||
outs[1] = data.candidates;
|
||||
ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TODO: Call backend_accept after all samplers have been applied.
|
||||
/*
|
||||
for (const auto & [seq_id, sampler] : samplers) {
|
||||
if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) {
|
||||
ggml_tensor * selected_token = it->second;
|
||||
if (selected_token != nullptr) {
|
||||
llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token);
|
||||
const auto it = sampling_rows.find(seq_id);
|
||||
if (it == sampling_rows.end()) {
|
||||
continue;
|
||||
}
|
||||
|
||||
for (uint32_t row : it->second) {
|
||||
ggml_tensor * selected_token = res->t_sampled[row];
|
||||
if (selected_token != nullptr && sampler->iface->backend_accept) {
|
||||
sampler->iface->backend_accept(sampler, ctx0, gf, selected_token);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+4
-4
@@ -904,10 +904,10 @@ public:
|
||||
|
||||
std::vector<ggml_tensor *> t_layer_inp;
|
||||
|
||||
std::map<llama_seq_id, ggml_tensor *> t_sampled_logits;
|
||||
std::map<llama_seq_id, ggml_tensor *> t_candidates;
|
||||
std::map<llama_seq_id, ggml_tensor *> t_sampled;
|
||||
std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
|
||||
std::vector<ggml_tensor *> t_sampled;
|
||||
std::vector<ggml_tensor *> t_sampled_probs;
|
||||
std::vector<ggml_tensor *> t_sampled_logits;
|
||||
std::vector<ggml_tensor *> t_candidates;
|
||||
|
||||
std::vector<llm_graph_input_ptr> inputs;
|
||||
std::vector<llm_graph_fused_node> fused_nodes;
|
||||
|
||||
+190
-83
@@ -467,9 +467,11 @@ static void llama_sampler_empty_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_empty_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
GGML_UNUSED(smpl);
|
||||
GGML_UNUSED(buft);
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
return true;
|
||||
}
|
||||
@@ -510,6 +512,7 @@ static struct llama_sampler_i llama_sampler_empty_i = {
|
||||
/* .backend_init = */ llama_sampler_empty_backend_init,
|
||||
/* .backend_accept = */ llama_sampler_empty_backend_accept,
|
||||
/* .backend_apply = */ llama_sampler_empty_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ llama_sampler_empty_backend_set_input,
|
||||
};
|
||||
|
||||
@@ -559,6 +562,60 @@ private:
|
||||
bool support;
|
||||
};
|
||||
|
||||
struct llama_sampler_backend_probe {
|
||||
ggml_context_ptr ctx;
|
||||
ggml_cgraph * gf;
|
||||
};
|
||||
|
||||
static llama_sampler_backend_probe llama_sampler_backend_probe_graph(
|
||||
llama_sampler * sampler,
|
||||
int64_t n_candidates,
|
||||
uint32_t max_nodes) {
|
||||
ggml_init_params params = {
|
||||
/*.mem_size =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),
|
||||
/*.mem_buffer =*/ nullptr,
|
||||
/*.no_alloc =*/ true,
|
||||
};
|
||||
|
||||
ggml_context_ptr ctx_ptr { ggml_init(params) };
|
||||
if (!ctx_ptr) {
|
||||
throw std::runtime_error(format("failed to create ggml context"));
|
||||
}
|
||||
|
||||
auto * ctx = ctx_ptr.get();
|
||||
auto * gf = ggml_new_graph_custom(ctx, max_nodes, false);
|
||||
|
||||
llama_sampler_data data = {
|
||||
/*.logits =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates),
|
||||
/*.probs =*/ nullptr,
|
||||
/*.sampled =*/ nullptr,
|
||||
/*.candidates =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates),
|
||||
};
|
||||
|
||||
if (sampler->iface->backend_reset) {
|
||||
sampler->iface->backend_reset(sampler);
|
||||
}
|
||||
sampler->iface->backend_apply(sampler, ctx, gf, &data);
|
||||
|
||||
for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {
|
||||
if (output) {
|
||||
ggml_build_forward_expand(gf, output);
|
||||
}
|
||||
}
|
||||
|
||||
return { std::move(ctx_ptr), gf };
|
||||
}
|
||||
|
||||
static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) {
|
||||
uint32_t n_tensors = 0;
|
||||
for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor;
|
||||
tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) {
|
||||
++n_tensors;
|
||||
}
|
||||
|
||||
return std::max<uint32_t>(ggml_graph_n_nodes(probe.gf), n_tensors);
|
||||
}
|
||||
|
||||
// check if all ggml ops used by the sampler are supported by the backend
|
||||
static bool llama_sampler_backend_support(
|
||||
llama_sampler * smpl,
|
||||
@@ -569,50 +626,10 @@ static bool llama_sampler_backend_support(
|
||||
return true;
|
||||
}
|
||||
|
||||
ggml_init_params params = {
|
||||
/*.mem_size =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
|
||||
/*.mem_buffer =*/ NULL,
|
||||
/*.no_alloc =*/ true,
|
||||
};
|
||||
auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE);
|
||||
|
||||
ggml_context_ptr ctx_ptr { ggml_init(params) };
|
||||
if (!ctx_ptr) {
|
||||
throw std::runtime_error(format("failed to create ggml context"));
|
||||
}
|
||||
|
||||
ggml_context * ctx = ctx_ptr.get();
|
||||
|
||||
const int64_t n = 1024*1024;
|
||||
|
||||
llama_sampler_data data = {
|
||||
/*.logits = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
|
||||
/*.probs = */ nullptr,
|
||||
/*.sampled = */ nullptr,
|
||||
/*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
|
||||
};
|
||||
|
||||
ggml_cgraph * gf = ggml_new_graph(ctx);
|
||||
|
||||
smpl->iface->backend_apply(smpl, ctx, gf, &data);
|
||||
|
||||
if (data.logits) {
|
||||
ggml_build_forward_expand(gf, data.logits);
|
||||
}
|
||||
|
||||
if (data.probs) {
|
||||
ggml_build_forward_expand(gf, data.probs);
|
||||
}
|
||||
|
||||
if (data.sampled) {
|
||||
ggml_build_forward_expand(gf, data.sampled);
|
||||
}
|
||||
|
||||
if (data.candidates) {
|
||||
ggml_build_forward_expand(gf, data.candidates);
|
||||
}
|
||||
|
||||
for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
|
||||
struct ggml_tensor * op = ggml_graph_node(gf, i);
|
||||
for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) {
|
||||
struct ggml_tensor * op = ggml_graph_node(probe.gf, i);
|
||||
|
||||
if (!ggml_backend_dev_supports_op(device, op)) {
|
||||
LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
|
||||
@@ -697,7 +714,8 @@ static void llama_sampler_chain_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_chain_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * chain = (llama_sampler_chain *) smpl->ctx;
|
||||
|
||||
GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
|
||||
@@ -705,15 +723,17 @@ static bool llama_sampler_chain_backend_init(
|
||||
chain->is_init = true;
|
||||
|
||||
bool res = true;
|
||||
bool backend_prefix = true;
|
||||
|
||||
for (auto & smpl : chain->samplers) {
|
||||
bool res_cur = true;
|
||||
bool res_cur = backend_prefix;
|
||||
|
||||
// to be able to run a sampler on the backend, it has to:
|
||||
// - have the .backend_init() API implemented
|
||||
// - return true during .backend_init()
|
||||
if (smpl.ptr->iface->backend_init) {
|
||||
if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
|
||||
// - support the requested output mode
|
||||
if (res_cur && smpl.ptr->iface->backend_init) {
|
||||
if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, require_multi_output)) {
|
||||
res_cur = false;
|
||||
}
|
||||
} else {
|
||||
@@ -721,10 +741,14 @@ static bool llama_sampler_chain_backend_init(
|
||||
}
|
||||
|
||||
smpl.is_backend = res_cur;
|
||||
backend_prefix = res_cur;
|
||||
|
||||
res = res && res_cur;
|
||||
}
|
||||
|
||||
auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE);
|
||||
chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
@@ -780,6 +804,19 @@ static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
|
||||
}
|
||||
}
|
||||
|
||||
static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) {
|
||||
auto * chain = (llama_sampler_chain *) smpl->ctx;
|
||||
|
||||
for (auto & entry : chain->samplers) {
|
||||
if (!entry.is_backend) {
|
||||
break;
|
||||
}
|
||||
if (entry.ptr->iface->backend_reset) {
|
||||
entry.ptr->iface->backend_reset(entry.ptr);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
static struct llama_sampler_i llama_sampler_chain_i = {
|
||||
/* .name = */ llama_sampler_chain_name,
|
||||
/* .accept = */ llama_sampler_chain_accept,
|
||||
@@ -790,6 +827,7 @@ static struct llama_sampler_i llama_sampler_chain_i = {
|
||||
/* .backend_init = */ llama_sampler_chain_backend_init,
|
||||
/* .backend_accept = */ llama_sampler_chain_backend_accept,
|
||||
/* .backend_apply = */ llama_sampler_chain_backend_apply,
|
||||
/* .backend_reset = */ llama_sampler_chain_backend_reset,
|
||||
/* .backend_set_input = */ llama_sampler_chain_backend_set_input,
|
||||
};
|
||||
|
||||
@@ -797,16 +835,27 @@ struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_param
|
||||
return llama_sampler_init(
|
||||
/* .iface = */ &llama_sampler_chain_i,
|
||||
/* .ctx = */ new llama_sampler_chain {
|
||||
/* .params = */ params,
|
||||
/* .is_init = */ false,
|
||||
/* .samplers = */ {},
|
||||
/* .cur = */ {},
|
||||
/* .t_sample_us = */ 0,
|
||||
/* .n_sample = */ 0,
|
||||
/* .params = */ params,
|
||||
/* .is_init = */ false,
|
||||
/* .n_nodes = */ 0,
|
||||
/* .samplers = */ {},
|
||||
/* .cur = */ {},
|
||||
/* .t_sample_us = */ 0,
|
||||
/* .n_sample = */ 0,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) {
|
||||
GGML_ASSERT(sampler != nullptr);
|
||||
GGML_ASSERT(sampler->iface == &llama_sampler_chain_i);
|
||||
|
||||
const auto * chain = (const llama_sampler_chain *) sampler->ctx;
|
||||
GGML_ASSERT(chain->is_init);
|
||||
|
||||
return chain->n_nodes;
|
||||
}
|
||||
|
||||
llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
|
||||
const llama_token sampled_token = llama_get_sampled_token_ith (ctx, idx);
|
||||
const float * sampled_probs = llama_get_sampled_probs_ith (ctx, idx);
|
||||
@@ -975,8 +1024,10 @@ static void llama_sampler_greedy_apply(struct llama_sampler * /*smpl*/, llama_to
|
||||
|
||||
static bool llama_sampler_greedy_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_greedy *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -1011,6 +1062,7 @@ static struct llama_sampler_i llama_sampler_greedy_i = {
|
||||
/* .backend_init = */ llama_sampler_greedy_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_greedy_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -1031,7 +1083,8 @@ struct llama_sampler_dist : public llama_sampler_backend {
|
||||
|
||||
std::mt19937 rng;
|
||||
|
||||
ggml_tensor * inp_uniform;
|
||||
// inputs for the current sampling graph
|
||||
std::vector<ggml_tensor *> inp_uniforms;
|
||||
};
|
||||
|
||||
static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
|
||||
@@ -1137,8 +1190,10 @@ static void llama_sampler_dist_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_dist_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_dist *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -1156,9 +1211,10 @@ static void llama_sampler_dist_backend_apply(
|
||||
|
||||
auto * sctx = (llama_sampler_dist *) smpl->ctx;
|
||||
|
||||
sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
|
||||
ggml_set_name (sctx->inp_uniform, "uniform");
|
||||
ggml_set_input(sctx->inp_uniform);
|
||||
ggml_tensor * inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
|
||||
ggml_format_name(inp_uniform, "uniform_%zu", sctx->inp_uniforms.size());
|
||||
ggml_set_input(inp_uniform);
|
||||
sctx->inp_uniforms.push_back(inp_uniform);
|
||||
|
||||
// flatten
|
||||
struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
|
||||
@@ -1174,7 +1230,7 @@ static void llama_sampler_dist_backend_apply(
|
||||
// Recall that each entry in cumsum is the cumulative probability up to that
|
||||
// index so values stay negative while the cumulative total is below the
|
||||
// random value, and become zero/positive once the threshold is crossed.
|
||||
struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
|
||||
struct ggml_tensor * diff = ggml_sub(ctx, cumsum, inp_uniform);
|
||||
ggml_set_name(diff, "dist_cumsum");
|
||||
|
||||
// The ggml_step function produces a tensor where entries are 1 if the
|
||||
@@ -1201,6 +1257,10 @@ static void llama_sampler_dist_backend_apply(
|
||||
|
||||
sampled_token = ggml_get_rows(ctx, candidates, idx);
|
||||
ggml_set_name(sampled_token, "dist_sampled_token");
|
||||
|
||||
// candidates may be a view whose backing storage can be reused
|
||||
data->candidates = ggml_cont(ctx, data->candidates);
|
||||
ggml_set_name(data->candidates, "dist_candidates_out");
|
||||
}
|
||||
|
||||
data->sampled = sampled_token;
|
||||
@@ -1210,7 +1270,7 @@ static void llama_sampler_dist_backend_apply(
|
||||
static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
|
||||
auto * sctx = (llama_sampler_dist *) smpl->ctx;
|
||||
|
||||
GGML_ASSERT(sctx->inp_uniform != nullptr);
|
||||
GGML_ASSERT(!sctx->inp_uniforms.empty());
|
||||
|
||||
// We sample in double precision and cast to float to match rnd numbers of
|
||||
// llama_dampler_dist which uses double precision (sampling from
|
||||
@@ -1218,9 +1278,17 @@ static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
|
||||
// std::uniform_real_distribution<float> with same rng will produce
|
||||
// different sequences).
|
||||
std::uniform_real_distribution<double> dist(0.0f, 1.0f);
|
||||
const float rnd = dist(sctx->rng);
|
||||
|
||||
ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
|
||||
for (auto * inp_uniform : sctx->inp_uniforms) {
|
||||
GGML_ASSERT(inp_uniform != nullptr);
|
||||
const float rnd = dist(sctx->rng);
|
||||
ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float));
|
||||
}
|
||||
}
|
||||
|
||||
static void llama_sampler_dist_backend_reset(struct llama_sampler * smpl) {
|
||||
auto * sctx = (llama_sampler_dist *) smpl->ctx;
|
||||
sctx->inp_uniforms.clear();
|
||||
}
|
||||
|
||||
static struct llama_sampler_i llama_sampler_dist_i = {
|
||||
@@ -1233,6 +1301,7 @@ static struct llama_sampler_i llama_sampler_dist_i = {
|
||||
/* .backend_init = */ llama_sampler_dist_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_dist_backend_apply,
|
||||
/* .backend_reset = */ llama_sampler_dist_backend_reset,
|
||||
/* .backend_set_input = */ llama_sampler_dist_backend_set_input,
|
||||
};
|
||||
|
||||
@@ -1242,10 +1311,10 @@ struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
|
||||
/* .iface = */ &llama_sampler_dist_i,
|
||||
/* .ctx = */ new llama_sampler_dist {
|
||||
("dist"),
|
||||
/* .seed = */ seed,
|
||||
/* .seed_cur = */ seed_cur,
|
||||
/* .rng = */ std::mt19937(seed_cur),
|
||||
/* .inp_uniform = */ nullptr,
|
||||
/* .seed = */ seed,
|
||||
/* .seed_cur = */ seed_cur,
|
||||
/* .rng = */ std::mt19937(seed_cur),
|
||||
/* .inp_uniforms = */ {},
|
||||
}
|
||||
);
|
||||
}
|
||||
@@ -1277,8 +1346,10 @@ static void llama_sampler_top_k_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_top_k_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_top_k *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -1324,6 +1395,7 @@ static struct llama_sampler_i llama_sampler_top_k_i = {
|
||||
/* .backend_init = */ llama_sampler_top_k_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_top_k_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -1423,8 +1495,10 @@ static void llama_sampler_top_p_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_top_p_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_top_p *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -1520,6 +1594,7 @@ static struct llama_sampler_i llama_sampler_top_p_i = {
|
||||
/* .backend_init = */ llama_sampler_top_p_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_top_p_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -1618,8 +1693,10 @@ static void llama_sampler_min_p_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_min_p_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_min_p *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -1679,6 +1756,7 @@ static struct llama_sampler_i llama_sampler_min_p_i = {
|
||||
/* .backend_init = */ llama_sampler_min_p_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_min_p_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -1789,6 +1867,7 @@ static struct llama_sampler_i llama_sampler_typical_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -1866,8 +1945,10 @@ static void llama_sampler_backend_temp_sampling(
|
||||
|
||||
static bool llama_sampler_temp_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_temp *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -1895,6 +1976,7 @@ static struct llama_sampler_i llama_sampler_temp_i = {
|
||||
/* .backend_init = */ llama_sampler_temp_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_temp_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -2009,8 +2091,10 @@ static void llama_sampler_temp_ext_free(struct llama_sampler * smpl) {
|
||||
|
||||
static bool llama_sampler_temp_ext_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
const bool res = llama_sampler_backend_support(smpl, buft);
|
||||
|
||||
@@ -2094,6 +2178,7 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = {
|
||||
/* .backend_init = */ llama_sampler_temp_ext_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_temp_ext_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -2201,6 +2286,7 @@ static struct llama_sampler_i llama_sampler_xtc_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -2320,6 +2406,7 @@ static struct llama_sampler_i llama_sampler_mirostat_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -2424,6 +2511,7 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -2545,6 +2633,7 @@ static struct llama_sampler_i llama_sampler_grammar_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -2962,6 +3051,7 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
|
||||
/* .backend_init = */ llama_sampler_penalties_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_penalties_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
|
||||
};
|
||||
|
||||
@@ -3057,6 +3147,7 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -3394,6 +3485,7 @@ static struct llama_sampler_i llama_sampler_dry_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -3613,6 +3705,7 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = {
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
@@ -3715,13 +3808,17 @@ static void llama_sampler_logit_bias_backend_apply(
|
||||
|
||||
const size_t n = sctx->logit_bias.size();
|
||||
|
||||
sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
|
||||
ggml_set_name(sctx->inp_logit_bias, "logit_bias");
|
||||
ggml_set_input(sctx->inp_logit_bias);
|
||||
if (sctx->inp_logit_bias == nullptr) {
|
||||
GGML_ASSERT(sctx->inp_logit_idxs == nullptr);
|
||||
|
||||
sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
|
||||
ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
|
||||
ggml_set_input(sctx->inp_logit_idxs);
|
||||
sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
|
||||
ggml_set_name(sctx->inp_logit_bias, "logit_bias");
|
||||
ggml_set_input(sctx->inp_logit_bias);
|
||||
|
||||
sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
|
||||
ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
|
||||
ggml_set_input(sctx->inp_logit_idxs);
|
||||
}
|
||||
|
||||
ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);
|
||||
|
||||
@@ -3756,10 +3853,18 @@ static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * sm
|
||||
ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
|
||||
}
|
||||
|
||||
static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) {
|
||||
auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
|
||||
sctx->inp_logit_bias = nullptr;
|
||||
sctx->inp_logit_idxs = nullptr;
|
||||
}
|
||||
|
||||
static bool llama_sampler_logit_bias_backend_init(
|
||||
struct llama_sampler * smpl,
|
||||
ggml_backend_buffer_type_t buft) {
|
||||
ggml_backend_buffer_type_t buft,
|
||||
bool require_multi_output) {
|
||||
GGML_UNUSED(buft);
|
||||
GGML_UNUSED(require_multi_output);
|
||||
|
||||
auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
|
||||
|
||||
@@ -3782,6 +3887,7 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = {
|
||||
/* .backend_init = */ llama_sampler_logit_bias_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ llama_sampler_logit_bias_backend_apply,
|
||||
/* .backend_reset = */ llama_sampler_logit_bias_backend_reset,
|
||||
/* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
|
||||
};
|
||||
|
||||
@@ -4022,10 +4128,11 @@ static struct llama_sampler_i llama_sampler_infill_i = {
|
||||
/* .reset = */ nullptr,
|
||||
/* .clone = */ llama_sampler_infill_clone,
|
||||
/* .free = */ llama_sampler_infill_free,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
/* .backend_init = */ nullptr,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ nullptr,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
|
||||
|
||||
@@ -15,6 +15,8 @@ struct llama_sampler_chain {
|
||||
// has .backend_init() been called?
|
||||
bool is_init = false;
|
||||
|
||||
uint32_t n_nodes = 0;
|
||||
|
||||
struct info {
|
||||
bool is_backend;
|
||||
|
||||
@@ -33,6 +35,8 @@ struct llama_sampler_chain {
|
||||
mutable int32_t n_sample;
|
||||
};
|
||||
|
||||
uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler);
|
||||
|
||||
struct llama_sampler * llama_sampler_init_dry_testing(
|
||||
float dry_multiplier,
|
||||
float dry_base,
|
||||
|
||||
@@ -2,7 +2,9 @@
|
||||
#include "common.h"
|
||||
#include "download.h"
|
||||
#include "llama.h"
|
||||
#include "speculative.h"
|
||||
|
||||
#include <limits>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <sstream>
|
||||
@@ -14,6 +16,14 @@
|
||||
static void test(void) {
|
||||
common_params params;
|
||||
|
||||
assert(common_speculative_n_outputs_max(16, 2, 3) == 8);
|
||||
assert(common_speculative_n_outputs_max(16, 2, -1) == 2);
|
||||
assert(common_speculative_n_outputs_max(4, 2, 3) == 4);
|
||||
assert(common_speculative_n_outputs_max(
|
||||
std::numeric_limits<int32_t>::max(),
|
||||
std::numeric_limits<int32_t>::max(),
|
||||
std::numeric_limits<int32_t>::max()) == std::numeric_limits<int32_t>::max());
|
||||
|
||||
printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n");
|
||||
for (int ex = 0; ex < LLAMA_EXAMPLE_COUNT; ex++) {
|
||||
try {
|
||||
|
||||
+332
-33
@@ -14,6 +14,7 @@
|
||||
#include <fstream>
|
||||
#include <functional>
|
||||
#include <map>
|
||||
#include <random>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
#include <unordered_set>
|
||||
@@ -80,7 +81,12 @@ struct test_context {
|
||||
std::unordered_map<llama_seq_id, int32_t> seq_positions;
|
||||
std::unordered_map<llama_seq_id, int32_t> last_batch_info;
|
||||
|
||||
test_context(const test_params & params, std::vector<llama_sampler_seq_config> & configs, int32_t n_seq_max = -1) {
|
||||
test_context(
|
||||
const test_params & params,
|
||||
std::vector<llama_sampler_seq_config> & configs,
|
||||
int32_t n_seq_max = -1,
|
||||
uint32_t n_outputs_max = 0,
|
||||
uint32_t n_ubatch = 0) {
|
||||
auto * model = params.model.get();
|
||||
|
||||
GGML_ASSERT(model);
|
||||
@@ -89,6 +95,10 @@ struct test_context {
|
||||
llama_context_params cparams = llama_context_default_params();
|
||||
cparams.n_ctx = 512;
|
||||
cparams.n_batch = 512;
|
||||
if (n_ubatch > 0) {
|
||||
cparams.n_ubatch = n_ubatch;
|
||||
}
|
||||
cparams.n_outputs_max = n_outputs_max;
|
||||
cparams.samplers = configs.data();
|
||||
cparams.n_samplers = configs.size();
|
||||
cparams.kv_unified = true;
|
||||
@@ -262,6 +272,65 @@ struct test_context {
|
||||
}
|
||||
};
|
||||
|
||||
struct test_single_output_backend_sampler {
|
||||
bool backend_initialized = false;
|
||||
bool backend_require_multi_output = false;
|
||||
int backend_apply_count = 0;
|
||||
int apply_count = 0;
|
||||
};
|
||||
|
||||
static const char * test_single_output_backend_sampler_name(const llama_sampler * /*smpl*/) {
|
||||
return "single-output-backend";
|
||||
}
|
||||
|
||||
static void test_single_output_backend_sampler_apply(
|
||||
llama_sampler * smpl, llama_token_data_array * /*cur_p*/) {
|
||||
auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
|
||||
ctx->apply_count++;
|
||||
}
|
||||
|
||||
static void test_single_output_backend_sampler_free(llama_sampler * smpl) {
|
||||
delete (test_single_output_backend_sampler *) smpl->ctx;
|
||||
}
|
||||
|
||||
static bool test_single_output_backend_sampler_backend_init(
|
||||
llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, bool require_multi_output) {
|
||||
auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
|
||||
ctx->backend_require_multi_output = require_multi_output;
|
||||
if (require_multi_output) {
|
||||
return false;
|
||||
}
|
||||
ctx->backend_initialized = true;
|
||||
return true;
|
||||
}
|
||||
|
||||
static void test_single_output_backend_sampler_backend_apply(
|
||||
llama_sampler * smpl, ggml_context * /*ctx*/, ggml_cgraph * /*gf*/, llama_sampler_data * /*data*/) {
|
||||
auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
|
||||
ctx->backend_apply_count++;
|
||||
}
|
||||
|
||||
static llama_sampler_i test_single_output_backend_sampler_i = {
|
||||
/* .name = */ test_single_output_backend_sampler_name,
|
||||
/* .accept = */ nullptr,
|
||||
/* .apply = */ test_single_output_backend_sampler_apply,
|
||||
/* .reset = */ nullptr,
|
||||
/* .clone = */ nullptr,
|
||||
/* .free = */ test_single_output_backend_sampler_free,
|
||||
/* .backend_init = */ test_single_output_backend_sampler_backend_init,
|
||||
/* .backend_accept = */ nullptr,
|
||||
/* .backend_apply = */ test_single_output_backend_sampler_backend_apply,
|
||||
/* .backend_reset = */ nullptr,
|
||||
/* .backend_set_input = */ nullptr,
|
||||
};
|
||||
|
||||
static llama_sampler * test_single_output_backend_sampler_init(
|
||||
test_single_output_backend_sampler ** sampler_ctx) {
|
||||
auto * ctx = new test_single_output_backend_sampler;
|
||||
*sampler_ctx = ctx;
|
||||
return llama_sampler_init(&test_single_output_backend_sampler_i, ctx);
|
||||
}
|
||||
|
||||
static void test_backend_greedy_sampling(const test_params & params) {
|
||||
const int seq_id = 0;
|
||||
|
||||
@@ -1527,43 +1596,270 @@ static void test_backend_cpu_mixed_batch(const test_params & params) {
|
||||
printf("backend-cpu mixed batch test PASSED\n");
|
||||
}
|
||||
|
||||
static void test_backend_max_outputs(const test_params & params) {
|
||||
const int seq_id = 0;
|
||||
const int32_t seed = 88;
|
||||
static void test_backend_multi_output_disabled(const test_params & params) {
|
||||
const llama_seq_id seq_id = 0;
|
||||
|
||||
llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
|
||||
llama_sampler_ptr backend_sampler_chain(llama_sampler_chain_init(backend_chain_params));
|
||||
llama_sampler_chain_add(backend_sampler_chain.get(), llama_sampler_init_dist(seed));
|
||||
std::vector<llama_sampler_seq_config> backend_sampler_configs = {{ seq_id, backend_sampler_chain.get() }};
|
||||
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88));
|
||||
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
|
||||
test_context test_ctx(params, configs, 1, 1);
|
||||
|
||||
test_context test_ctx(params, backend_sampler_configs);
|
||||
llama_batch batch = llama_batch_init(2, 0, 1);
|
||||
common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), 0, { seq_id }, true);
|
||||
common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), 1, { seq_id }, true);
|
||||
|
||||
llama_batch batch = llama_batch_init(512, 0, 1);
|
||||
std::string prompt = "Hello";
|
||||
|
||||
std::vector<llama_token> tokens;
|
||||
tokens.push_back(llama_vocab_bos(test_ctx.vocab));
|
||||
|
||||
std::vector<llama_token> prompt_tokens(32);
|
||||
int n_tokens = llama_tokenize(test_ctx.vocab, prompt.c_str(), prompt.length(),
|
||||
prompt_tokens.data(), prompt_tokens.size(),
|
||||
false, false);
|
||||
for (int i = 0; i < n_tokens; i++) {
|
||||
tokens.push_back(prompt_tokens[i]);
|
||||
}
|
||||
|
||||
for (size_t i = 0; i < tokens.size(); i++) {
|
||||
// set all tokens as output to trigger error
|
||||
common_batch_add(batch, tokens[i], i, { seq_id }, true);
|
||||
}
|
||||
|
||||
printf(">>> test_max_outputs expected error start:\n");
|
||||
printf(">>> test_backend_multi_output_disabled expected error start:\n");
|
||||
const int ret = llama_decode(test_ctx.ctx.get(), batch);
|
||||
GGML_ASSERT(ret != 0 && "llama_decode should not succeed multiple outputs per sequence");
|
||||
printf("<<< test_max_outputs expected error end.\n");
|
||||
GGML_ASSERT(ret != 0 && "llama_decode should reject multiple outputs for one sequence");
|
||||
printf("<<< test_backend_multi_output_disabled expected error end.\n");
|
||||
|
||||
llama_batch_free(batch);
|
||||
|
||||
printf("backend max outputs test PASSED\n");
|
||||
printf("backend multi-output disabled test PASSED\n");
|
||||
}
|
||||
|
||||
// greedy is a stateless terminal selector; verify multi-output backend argmax
|
||||
// matches the per-row argmax of the reference logits.
|
||||
static void test_backend_multi_output_greedy(const test_params & params) {
|
||||
const llama_seq_id seq_id = 0;
|
||||
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
|
||||
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
|
||||
|
||||
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_greedy());
|
||||
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
|
||||
test_context test_ctx(params, configs, 1, 4);
|
||||
|
||||
std::vector<llama_sampler_seq_config> reference_configs;
|
||||
test_context reference_ctx(params, reference_configs, 1, 4);
|
||||
|
||||
llama_batch batch = llama_batch_init(4, 0, 1);
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
common_batch_add(batch, llama_vocab_bos(vocab), i, { seq_id }, true);
|
||||
}
|
||||
|
||||
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0 &&
|
||||
"multi-output backend greedy sampling should succeed");
|
||||
GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
|
||||
|
||||
for (int i = 0; i < batch.n_tokens; ++i) {
|
||||
const llama_token backend_token = llama_get_sampled_token_ith(test_ctx.ctx.get(), i);
|
||||
GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
|
||||
|
||||
const float * logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
|
||||
GGML_ASSERT(logits != nullptr);
|
||||
llama_token argmax = 0;
|
||||
for (llama_token t = 1; t < n_vocab; ++t) {
|
||||
if (logits[t] > logits[argmax]) {
|
||||
argmax = t;
|
||||
}
|
||||
}
|
||||
printf("row %d: backend greedy=%d argmax=%d\n", i, backend_token, argmax);
|
||||
GGML_ASSERT(backend_token == argmax);
|
||||
}
|
||||
|
||||
llama_batch_free(batch);
|
||||
|
||||
printf("backend multi-output greedy test PASSED\n");
|
||||
}
|
||||
|
||||
static void test_backend_multi_output_sampling_chain(const test_params & params) {
|
||||
const llama_seq_id seq_id = 0;
|
||||
const int32_t seed = 88;
|
||||
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
|
||||
const int32_t n_vocab = llama_vocab_n_tokens(vocab);
|
||||
const uint32_t k = std::min<uint32_t>(2048, n_vocab);
|
||||
|
||||
const llama_logit_bias bias = { llama_vocab_bos(vocab), -0.1f };
|
||||
|
||||
auto make_filter_chain = [&]() {
|
||||
llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
|
||||
llama_sampler_chain_add(result.get(), llama_sampler_init_logit_bias(
|
||||
n_vocab, 1, &bias));
|
||||
llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
|
||||
llama_sampler_chain_add(result.get(), llama_sampler_init_top_p(0.9f, 1));
|
||||
llama_sampler_chain_add(result.get(), llama_sampler_init_min_p(0.01f, 1));
|
||||
llama_sampler_chain_add(result.get(), llama_sampler_init_temp(0.8f));
|
||||
return result;
|
||||
};
|
||||
|
||||
llama_sampler_ptr chain = make_filter_chain();
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
|
||||
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
|
||||
test_context test_ctx(params, configs, 1, 4, 2);
|
||||
|
||||
std::vector<llama_sampler_seq_config> reference_configs;
|
||||
test_context reference_ctx(params, reference_configs, 1, 4, 2);
|
||||
llama_sampler_ptr reference_filters = make_filter_chain();
|
||||
|
||||
std::vector<llama_token_data> reference_data(n_vocab);
|
||||
std::mt19937 reference_rng(seed);
|
||||
std::uniform_real_distribution<double> reference_dist(0.0, 1.0);
|
||||
int32_t n_reused_after_first_round = -1;
|
||||
|
||||
for (int round = 0; round < 2; ++round) {
|
||||
llama_batch batch = llama_batch_init(4, 0, 1);
|
||||
for (int i = 0; i < 4; ++i) {
|
||||
common_batch_add(batch, llama_vocab_bos(vocab), round * 4 + i, { seq_id }, true);
|
||||
}
|
||||
|
||||
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
|
||||
GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
|
||||
|
||||
for (int i = 0; i < batch.n_tokens; ++i) {
|
||||
const llama_token backend_token = llama_get_sampled_token_ith(test_ctx.ctx.get(), i);
|
||||
const float * sampled_logits = llama_get_sampled_logits_ith(test_ctx.ctx.get(), i);
|
||||
const float * sampled_probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), i);
|
||||
const llama_token * sampled_candidates = llama_get_sampled_candidates_ith(test_ctx.ctx.get(), i);
|
||||
const uint32_t n_logits = llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i);
|
||||
const uint32_t n_probs = llama_get_sampled_probs_count_ith(test_ctx.ctx.get(), i);
|
||||
const uint32_t n_candidates = llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i);
|
||||
|
||||
GGML_ASSERT(backend_token >= 0 && backend_token < n_vocab);
|
||||
GGML_ASSERT(sampled_logits != nullptr);
|
||||
GGML_ASSERT(sampled_probs != nullptr);
|
||||
GGML_ASSERT(sampled_candidates != nullptr);
|
||||
GGML_ASSERT(n_logits == k);
|
||||
GGML_ASSERT(n_probs == n_logits);
|
||||
GGML_ASSERT(n_candidates == n_logits);
|
||||
|
||||
const float * reference_logits = llama_get_logits_ith(reference_ctx.ctx.get(), i);
|
||||
GGML_ASSERT(reference_logits != nullptr);
|
||||
|
||||
for (llama_token token = 0; token < n_vocab; ++token) {
|
||||
reference_data[token] = { token, reference_logits[token], 0.0f };
|
||||
}
|
||||
|
||||
llama_token_data_array reference = {
|
||||
/* .data = */ reference_data.data(),
|
||||
/* .size = */ reference_data.size(),
|
||||
/* .selected = */ LLAMA_TOKEN_NULL,
|
||||
/* .sorted = */ false,
|
||||
};
|
||||
llama_sampler_apply(reference_filters.get(), &reference);
|
||||
|
||||
std::vector<bool> reference_matched(reference.size, false);
|
||||
const llama_token_data * reference_begin = reference.data;
|
||||
const llama_token_data * reference_end = reference.data + reference.size;
|
||||
size_t n_matched = 0;
|
||||
int32_t sampled_index = -1;
|
||||
float prob_sum = 0.0f;
|
||||
|
||||
for (uint32_t j = 0; j < n_logits; ++j) {
|
||||
GGML_ASSERT(std::isfinite(sampled_probs[j]));
|
||||
GGML_ASSERT(sampled_probs[j] >= 0.0f);
|
||||
prob_sum += sampled_probs[j];
|
||||
|
||||
if (sampled_candidates[j] == backend_token) {
|
||||
sampled_index = j;
|
||||
}
|
||||
|
||||
if (!std::isfinite(sampled_logits[j])) {
|
||||
GGML_ASSERT(std::isinf(sampled_logits[j]) && sampled_logits[j] < 0.0f);
|
||||
GGML_ASSERT(sampled_probs[j] == 0.0f);
|
||||
continue;
|
||||
}
|
||||
|
||||
const llama_token_data * match = std::find_if(reference_begin, reference_end,
|
||||
[&](const llama_token_data & candidate) {
|
||||
return candidate.id == sampled_candidates[j];
|
||||
});
|
||||
GGML_ASSERT(match != reference_end);
|
||||
|
||||
const size_t i_reference = match - reference.data;
|
||||
GGML_ASSERT(!reference_matched[i_reference]);
|
||||
|
||||
const float tolerance = 1e-4f * std::max(1.0f, std::fabs(match->logit));
|
||||
GGML_ASSERT(std::fabs(sampled_logits[j] - match->logit) <= tolerance);
|
||||
reference_matched[i_reference] = true;
|
||||
++n_matched;
|
||||
}
|
||||
|
||||
GGML_ASSERT(n_matched == reference.size);
|
||||
GGML_ASSERT(sampled_index >= 0);
|
||||
GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
|
||||
|
||||
const float rnd = reference_dist(reference_rng);
|
||||
float cumsum_before = 0.0f;
|
||||
for (int32_t j = 0; j < sampled_index; ++j) {
|
||||
cumsum_before += sampled_probs[j];
|
||||
}
|
||||
const float cumsum_sampled = cumsum_before + sampled_probs[sampled_index];
|
||||
GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
|
||||
GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
|
||||
}
|
||||
|
||||
llama_batch_free(batch);
|
||||
|
||||
const int32_t n_reused = llama_perf_context(test_ctx.ctx.get()).n_reused;
|
||||
if (round == 0) {
|
||||
n_reused_after_first_round = n_reused;
|
||||
} else {
|
||||
GGML_ASSERT(n_reused > n_reused_after_first_round);
|
||||
}
|
||||
}
|
||||
|
||||
printf("backend multi-output sampling chain test PASSED\n");
|
||||
}
|
||||
|
||||
static void test_backend_multi_output_cpu_suffix(const test_params & params) {
|
||||
const llama_seq_id seq_id = 0;
|
||||
const int32_t k = 8;
|
||||
const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
|
||||
|
||||
{
|
||||
test_single_output_backend_sampler * sampler_ctx = nullptr;
|
||||
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_top_k(k));
|
||||
llama_sampler_chain_add(chain.get(), test_single_output_backend_sampler_init(&sampler_ctx));
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88));
|
||||
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
|
||||
test_context test_ctx(params, configs, 1, 1);
|
||||
|
||||
llama_batch batch = llama_batch_init(1, 0, 1);
|
||||
common_batch_add(batch, llama_vocab_bos(vocab), 0, { seq_id }, true);
|
||||
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
|
||||
|
||||
GGML_ASSERT(sampler_ctx->backend_initialized);
|
||||
GGML_ASSERT(!sampler_ctx->backend_require_multi_output);
|
||||
GGML_ASSERT(sampler_ctx->backend_apply_count > 0);
|
||||
GGML_ASSERT(sampler_ctx->apply_count == 0);
|
||||
GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), 0) != LLAMA_TOKEN_NULL);
|
||||
|
||||
llama_batch_free(batch);
|
||||
}
|
||||
|
||||
{
|
||||
test_single_output_backend_sampler * sampler_ctx = nullptr;
|
||||
llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_top_k(k));
|
||||
llama_sampler_chain_add(chain.get(), test_single_output_backend_sampler_init(&sampler_ctx));
|
||||
llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88));
|
||||
std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
|
||||
test_context test_ctx(params, configs, 1, 2);
|
||||
|
||||
llama_batch batch = llama_batch_init(2, 0, 1);
|
||||
for (int i = 0; i < 2; ++i) {
|
||||
common_batch_add(batch, llama_vocab_bos(vocab), i, { seq_id }, true);
|
||||
}
|
||||
GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
|
||||
|
||||
GGML_ASSERT(!sampler_ctx->backend_initialized);
|
||||
GGML_ASSERT(sampler_ctx->backend_require_multi_output);
|
||||
GGML_ASSERT(sampler_ctx->backend_apply_count == 0);
|
||||
for (int i = 0; i < batch.n_tokens; ++i) {
|
||||
GGML_ASSERT(llama_get_sampled_token_ith(test_ctx.ctx.get(), i) == LLAMA_TOKEN_NULL);
|
||||
GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
|
||||
GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
|
||||
const llama_token token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i);
|
||||
GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
|
||||
}
|
||||
GGML_ASSERT(sampler_ctx->apply_count == batch.n_tokens);
|
||||
|
||||
llama_batch_free(batch);
|
||||
}
|
||||
|
||||
printf("backend multi-output CPU suffix test PASSED\n");
|
||||
}
|
||||
|
||||
struct backend_test_case {
|
||||
@@ -1583,7 +1879,10 @@ static const backend_test_case BACKEND_TESTS[] = {
|
||||
{ "dist", test_backend_dist_sampling, true },
|
||||
{ "dist_and_cpu", test_backend_dist_sampling_and_cpu, true },
|
||||
{ "set_sampler", test_backend_set_sampler, true },
|
||||
{ "max_outputs", test_backend_max_outputs, true },
|
||||
{ "multi_output_disabled", test_backend_multi_output_disabled, true },
|
||||
{ "multi_output_greedy", test_backend_multi_output_greedy, true },
|
||||
{ "multi_output_sampling_chain", test_backend_multi_output_sampling_chain, true },
|
||||
{ "multi_output_cpu", test_backend_multi_output_cpu_suffix, true },
|
||||
{ "mixed", test_backend_mixed_sampling, true },
|
||||
{ "min_p", test_backend_min_p_sampling, true },
|
||||
{ "cpu_mixed", test_backend_cpu_mixed_batch, true },
|
||||
|
||||
@@ -40,18 +40,15 @@ using json = nlohmann::ordered_json;
|
||||
constexpr int HTTP_POLLING_SECONDS = 1;
|
||||
|
||||
static uint32_t server_n_outputs_max(const common_params & params) {
|
||||
const uint32_t n_batch = params.n_batch;
|
||||
|
||||
if (params.embedding ||
|
||||
(params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
|
||||
return n_batch;
|
||||
return params.n_batch;
|
||||
}
|
||||
|
||||
const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(¶ms.speculative);
|
||||
const int32_t n_outputs = common_speculative_n_outputs_max(
|
||||
params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative));
|
||||
|
||||
const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
|
||||
|
||||
return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
|
||||
return std::max<int32_t>(1, n_outputs);
|
||||
}
|
||||
|
||||
// state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
|
||||
@@ -303,6 +300,7 @@ struct server_slot {
|
||||
json json_schema;
|
||||
|
||||
common_sampler_ptr smpl;
|
||||
bool backend_sampling = false;
|
||||
|
||||
llama_token sampled; // in speculative mode, this is the last accepted token
|
||||
|
||||
@@ -364,6 +362,7 @@ struct server_slot {
|
||||
task.reset();
|
||||
|
||||
llama_set_sampler(ctx_tgt, id, nullptr);
|
||||
backend_sampling = false;
|
||||
|
||||
// clear alora start
|
||||
alora_invocation_start = -1;
|
||||
@@ -1832,21 +1831,17 @@ private:
|
||||
|
||||
const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs;
|
||||
|
||||
bool backend_sampling = true;
|
||||
|
||||
backend_sampling &= task.params.sampling.backend_sampling;
|
||||
|
||||
// TODO: speculative decoding requires multiple samples per batch - not supported yet
|
||||
backend_sampling &= !(slot.can_speculate());
|
||||
bool use_backend_sampling = task.params.sampling.backend_sampling;
|
||||
|
||||
// TODO: getting pre sampling logits is not yet supported with backend sampling
|
||||
backend_sampling &= !need_pre_sample_logits;
|
||||
use_backend_sampling &= !need_pre_sample_logits;
|
||||
|
||||
// TODO: tmp until backend sampling is fully implemented
|
||||
if (backend_sampling) {
|
||||
llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
|
||||
if (use_backend_sampling) {
|
||||
slot.backend_sampling = llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
|
||||
} else {
|
||||
llama_set_sampler(ctx_tgt, slot.id, nullptr);
|
||||
slot.backend_sampling = false;
|
||||
}
|
||||
|
||||
SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str());
|
||||
@@ -3915,7 +3910,16 @@ private:
|
||||
slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);
|
||||
|
||||
slot.prompt.tokens.keep_first(ckpt.n_tokens);
|
||||
const bool restore_backend_sampler = slot.backend_sampling;
|
||||
if (restore_backend_sampler) {
|
||||
llama_set_sampler(slot.ctx_tgt, slot.id, nullptr);
|
||||
}
|
||||
|
||||
slot.smpl = std::move(smpl_save);
|
||||
if (restore_backend_sampler) {
|
||||
slot.backend_sampling = llama_set_sampler(
|
||||
slot.ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
|
||||
}
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -40,6 +40,7 @@ def test_with_and_without_draft():
|
||||
|
||||
# create new server with draft model
|
||||
create_server()
|
||||
server.backend_sampling = True
|
||||
server.start()
|
||||
res = server.make_request("POST", "/completion", data={
|
||||
"prompt": "I believe the meaning of life is",
|
||||
|
||||
Reference in New Issue
Block a user