diff --git a/common/common.cpp b/common/common.cpp index 84d4cdaa86..2e3f14cd1c 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1639,7 +1639,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.n_seq_max = params.n_parallel; cparams.n_rs_seq = params.speculative.need_n_rs_seq(); cparams.n_outputs_max = std::max(params.n_outputs_max, 0); - cparams.n_sampling_outputs_per_seq_max = std::max(params.n_sampling_outputs_per_seq_max, 0); + cparams.n_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0); cparams.n_batch = params.n_batch; cparams.n_ubatch = params.n_ubatch; cparams.n_threads = params.cpuparams.n_threads; diff --git a/common/common.h b/common/common.h index 03255108c0..2e0c3a4cf8 100644 --- a/common/common.h +++ b/common/common.h @@ -447,7 +447,7 @@ struct common_params { int32_t n_parallel = 1; // number of parallel sequences to decode int32_t n_sequences = 1; // number of sequences to decode int32_t n_outputs_max = 0; // max outputs in a batch (0 = n_batch) - int32_t n_sampling_outputs_per_seq_max = 1; // max outputs per sequence with backend sampling + int32_t n_outputs_max_per_seq = 1; // max outputs per sequence int32_t grp_attn_n = 1; // group-attention factor int32_t grp_attn_w = 512; // group-attention width int32_t n_print = -1; // print token count every n tokens (-1 = disabled) diff --git a/common/llguidance.cpp b/common/llguidance.cpp index f72dbd1c1e..500bb09147 100644 --- a/common/llguidance.cpp +++ b/common/llguidance.cpp @@ -117,6 +117,7 @@ static llama_sampler_i llama_sampler_llg_i = { /* .backend_apply = */ NULL, /* .backend_set_input = */ NULL, /* .backend_reset = */ NULL, + /* .copy_state = */ NULL, }; static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len, diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 58ad2bc07f..4884299f30 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -218,6 +218,7 @@ static struct llama_sampler_i common_reasoning_budget_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) { diff --git a/common/sampling.cpp b/common/sampling.cpp index ec9c885ddf..06dea1e1cc 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -518,6 +518,26 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) { }; } +void common_sampler_copy(const common_sampler * src, common_sampler * dst) { + if (!src || !dst || src == dst) { + return; + } + + GGML_ASSERT((src->grmr == nullptr) == (dst->grmr == nullptr)); + GGML_ASSERT((src->rbudget == nullptr) == (dst->rbudget == nullptr)); + + llama_sampler_copy(src->grmr, dst->grmr); + llama_sampler_copy(src->rbudget, dst->rbudget); + llama_sampler_copy(src->chain, dst->chain); + + dst->params = src->params; + dst->prev = src->prev; + dst->cur = src->cur; + dst->cur_p = src->cur_p; + dst->cur_p.data = src->cur_p.data ? dst->cur.data() : nullptr; // re-point to dst's buffer + dst->t_total_us = src->t_total_us; +} + void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) { // TODO: measure grammar performance diff --git a/common/sampling.h b/common/sampling.h index cb90d4ae7a..ced3c8364b 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -47,6 +47,7 @@ void common_sampler_free(struct common_sampler * gsmpl); void common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated); void common_sampler_reset (struct common_sampler * gsmpl); struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl); +void common_sampler_copy (const struct common_sampler * src, struct common_sampler * dst); // arguments can be nullptr to skip printing void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl); diff --git a/common/speculative.cpp b/common/speculative.cpp index a074ab608b..a25e8e5a34 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2292,7 +2292,7 @@ common_params common_base_params_to_speculative(const common_params & params) { result.cache_type_k = params_spec.cache_type_k; result.cache_type_v = params_spec.cache_type_v; result.n_outputs_max = params.n_parallel; - result.n_sampling_outputs_per_seq_max = 1; + result.n_outputs_max_per_seq = 1; return result; } diff --git a/docs/speculative.md b/docs/speculative.md index 3957db85c9..25abef1b60 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -202,6 +202,12 @@ Example Video: If a draft model is combined with a draftless decoding the draftless decoding has higher precedence. +### Backend Sampling + +Use `--backend-sampling` to run supported target-model samplers on the model backend. Draft-model sampling uses the backend by default and can be controlled with `--spec-draft-backend-sampling` and `--no-spec-draft-backend-sampling`. + +Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required. + ### General Speculative Parameters ``` diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp index 75ee27c299..6621058655 100644 --- a/examples/lookup/lookup.cpp +++ b/examples/lookup/lookup.cpp @@ -31,7 +31,7 @@ int main(int argc, char ** argv){ const auto output_limits = common_speculative_get_output_limits(params.n_batch, params.n_parallel, n_draft); params.n_outputs_max = output_limits.total; - params.n_sampling_outputs_per_seq_max = output_limits.per_seq; + params.n_outputs_max_per_seq = output_limits.per_seq; // init llama.cpp llama_backend_init(); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index c30fe11531..c727e8139d 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -33,7 +33,7 @@ int main(int argc, char ** argv) { const auto output_limits = common_speculative_get_output_limits( params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative)); params.n_outputs_max = output_limits.total; - params.n_sampling_outputs_per_seq_max = output_limits.per_seq; + params.n_outputs_max_per_seq = output_limits.per_seq; // init llama.cpp llama_backend_init(); @@ -62,7 +62,7 @@ int main(int argc, char ** argv) { auto params_dft = params; params_dft.n_outputs_max = params.n_parallel; - params_dft.n_sampling_outputs_per_seq_max = 1; + params_dft.n_outputs_max_per_seq = 1; params_dft.devices = params_spec.devices; params_dft.model = params_spec.mparams; diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index b6838b0675..17071aa054 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -61,7 +61,7 @@ int main(int argc, char ** argv) { const auto output_limits = common_speculative_get_output_limits( params.n_batch, params.n_parallel, params.speculative.draft.n_max); params.n_outputs_max = output_limits.total; - params.n_sampling_outputs_per_seq_max = output_limits.per_seq; + params.n_outputs_max_per_seq = output_limits.per_seq; // probability threshold for splitting a draft branch (only for n_seq_dft > 1) const float p_draft_split = params.speculative.draft.p_split; @@ -90,7 +90,7 @@ int main(int argc, char ** argv) { params.model = params.speculative.draft.mparams; params.n_gpu_layers = params.speculative.draft.n_gpu_layers; params.n_outputs_max = params.n_parallel; - params.n_sampling_outputs_per_seq_max = 1; + params.n_outputs_max_per_seq = 1; if (params.speculative.draft.cpuparams.n_threads > 0) { params.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads; } diff --git a/include/llama.h b/include/llama.h index f33e9ee8c1..c14eb6f50d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -348,15 +348,15 @@ extern "C" { // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations // https://github.com/ggml-org/llama.cpp/pull/7544 struct llama_context_params { - uint32_t n_ctx; // text context, 0 = from model - uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode - uint32_t n_ubatch; // physical maximum batch size - uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) - uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] - uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) - uint32_t n_sampling_outputs_per_seq_max; // max outputs per sequence with backend sampling (0 = n_outputs_max) - int32_t n_threads; // number of threads to use for generation - int32_t n_threads_batch; // number of threads to use for batch processing + uint32_t n_ctx; // text context, 0 = from model + uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode + uint32_t n_ubatch; // physical maximum batch size + uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models) + uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback (0 = no rollback) [EXPERIMENTAL] + uint32_t n_outputs_max; // max outputs in a ubatch (0 = n_batch) + uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max) + int32_t n_threads; // number of threads to use for generation + int32_t n_threads_batch; // number of threads to use for batch processing enum llama_context_type ctx_type; // set the context type (e.g. MTP) enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` @@ -1274,12 +1274,12 @@ extern "C" { // [EXPERIMENTAL] // backend sampling interface: - // return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_per_seq_max outputs per sequence + // return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_max_per_seq outputs per sequence // note: call once per sampler bool (*backend_init)( struct llama_sampler * smpl, ggml_backend_buffer_type_t buft, - uint32_t n_outputs_per_seq_max); + uint32_t n_outputs_max_per_seq); // call after .backend_apply() void (*backend_accept)( @@ -1300,6 +1300,10 @@ extern "C" { // called before rebuilding a sampling graph to clear any internal sampler state void (*backend_reset)(struct llama_sampler * smpl); + + // copy mutable state from src into dst while keeping dst's references to the current sampling graph + // src and dst must have the same type and configuration + void (*copy_state)(const struct llama_sampler * src, struct llama_sampler * dst); }; struct llama_sampler { @@ -1320,6 +1324,9 @@ extern "C" { LLAMA_API void llama_sampler_apply ( struct llama_sampler * smpl, llama_token_data_array * cur_p); LLAMA_API void llama_sampler_reset ( struct llama_sampler * smpl); LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl); + // copy mutable sampler state without changing dst or its sampling graph bindings + // src and dst must have the same type and configuration + LLAMA_API void llama_sampler_copy (const struct llama_sampler * src, struct llama_sampler * dst); // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add) LLAMA_API void llama_sampler_free ( struct llama_sampler * smpl); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 61f4df06db..4d66a274b9 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -247,8 +247,8 @@ llama_context::llama_context( cparams.n_ubatch = std::min(cparams.n_batch, params.n_ubatch == 0 ? params.n_batch : params.n_ubatch); cparams.n_outputs_max = params.n_outputs_max == 0 || llama_model_has_encoder(&model) ? cparams.n_batch : params.n_outputs_max; - cparams.n_sampling_outputs_per_seq_max = params.n_sampling_outputs_per_seq_max == 0 ? - cparams.n_outputs_max : std::min(params.n_sampling_outputs_per_seq_max, cparams.n_outputs_max); + cparams.n_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ? + cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, cparams.n_outputs_max); // Initialize backend samplers here so they are part of the sampling graph // before the reserve passes run later in this function. This avoids a later @@ -303,19 +303,19 @@ llama_context::llama_context( } } - LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max); - LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); - LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq); - LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); - LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); - LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); - LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); - LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); - LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); - LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); - LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq); - LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); - LLAMA_LOG_INFO("%s: n_sampling_outputs_per_seq_max = %u\n", __func__, cparams.n_sampling_outputs_per_seq_max); + LLAMA_LOG_INFO("%s: n_seq_max = %u\n", __func__, cparams.n_seq_max); + LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); + LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq); + LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); + LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); + LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); + LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); + LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); + LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); + LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); + LLAMA_LOG_INFO("%s: n_rs_seq = %u\n", __func__, cparams.n_rs_seq); + LLAMA_LOG_INFO("%s: n_outputs_max = %u\n", __func__, cparams.n_outputs_max); + LLAMA_LOG_INFO("%s: n_outputs_max_per_seq = %u\n", __func__, cparams.n_outputs_max_per_seq); if (cparams.n_ctx_seq < hparams.n_ctx_train) { LLAMA_LOG_INFO("%s: n_ctx_seq (%u) < n_ctx_train (%u) -- the full capacity of the model will not be utilized\n", @@ -1235,7 +1235,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, cparams.n_sampling_outputs_per_seq_max); + sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq); sampling.samplers[seq_id] = sampler; @@ -1681,9 +1681,9 @@ int llama_context::decode(const llama_batch & batch_inp) { seq_output_count[seq_id]++; auto sampler = sampling.samplers.find(seq_id); if (sampler != sampling.samplers.end() && - seq_output_count[seq_id] > (int32_t) cparams.n_sampling_outputs_per_seq_max) { + seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) { LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence " - "(seq_id %d had %d)\n", __func__, cparams.n_sampling_outputs_per_seq_max, + "(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq, seq_id, seq_output_count[seq_id]); return -1; } @@ -2314,14 +2314,14 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { 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.n_sampling_outputs_per_seq_max > 1) { + if (cparams.n_outputs_max_per_seq > 1) { n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes); } } const uint32_t n_sampling_outputs_max = std::min( std::min(n_tokens, cparams.n_outputs_max), - (uint64_t) cparams.n_seq_max * cparams.n_sampling_outputs_per_seq_max); + (uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq); res += n_sampling_nodes; if (n_sampling_outputs_max > 1) { @@ -2384,7 +2384,7 @@ ggml_cgraph * llama_context::graph_reserve( } const uint32_t n_sampling_outputs_per_seq = std::min( - ubatch.n_seq_tokens, cparams.n_sampling_outputs_per_seq_max); + ubatch.n_seq_tokens, cparams.n_outputs_max_per_seq); // select sampling rows in round-robin order across sampler sequences if (!sampler_seqs.empty()) { @@ -3496,7 +3496,7 @@ llama_context_params llama_context_default_params() { /*.n_seq_max =*/ 1, /*.n_rs_seq =*/ 0, /*.n_outputs_max =*/ 0, - /*.n_sampling_outputs_per_seq_max =*/ 1, + /*.n_outputs_max_per_seq =*/ 1, /*.n_threads =*/ GGML_DEFAULT_N_THREADS, // TODO: better default /*.n_threads_batch =*/ GGML_DEFAULT_N_THREADS, /*.ctx_type =*/ LLAMA_CONTEXT_TYPE_DEFAULT, diff --git a/src/llama-cparams.h b/src/llama-cparams.h index dbd4021c9f..574ce95920 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -15,7 +15,7 @@ struct llama_cparams { uint32_t n_seq_max; uint32_t n_rs_seq; // number of recurrent-state snapshots per seq for rollback uint32_t n_outputs_max; // max outputs supported by the context - uint32_t n_sampling_outputs_per_seq_max; + uint32_t n_outputs_max_per_seq; int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 37c6eb6d50..eb8d65b8a0 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -468,10 +468,10 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { GGML_UNUSED(smpl); GGML_UNUSED(buft); - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); return true; } @@ -514,6 +514,7 @@ static struct llama_sampler_i llama_sampler_empty_i = { /* .backend_apply = */ llama_sampler_empty_backend_apply, /* .backend_set_input = */ llama_sampler_empty_backend_set_input, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_empty(const char * name) { @@ -554,6 +555,12 @@ struct llama_sampler_backend { this->support = support; } + // copy the state that is not tied to the current sampling graph + // samplers that hold only immutable configuration can use this as is + void copy_state(const llama_sampler_backend & src) { + GGML_UNUSED(src); + } + private: std::string name; std::string name_ext; @@ -562,6 +569,12 @@ private: bool support; }; +// .copy_state for samplers deriving from llama_sampler_backend +template +static void llama_sampler_backend_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) { + ((T *) dst->ctx)->copy_state(*(const T *) src->ctx); +} + struct llama_sampler_backend_probe { ggml_context_ptr ctx; ggml_cgraph * gf; @@ -720,7 +733,7 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * chain = (llama_sampler_chain *) smpl->ctx; GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice"); @@ -731,24 +744,24 @@ static bool llama_sampler_chain_backend_init( bool backend_prefix = true; for (auto & smpl : chain->samplers) { - bool res_cur = backend_prefix; + bool cur_prefix = 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() // - support the requested per-sequence output limit - if (res_cur && smpl.ptr->iface->backend_init) { - if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_per_seq_max)) { - res_cur = false; + if (cur_prefix && smpl.ptr->iface->backend_init) { + if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) { + cur_prefix = false; } } else { - res_cur = false; + cur_prefix = false; } - smpl.is_backend = res_cur; - backend_prefix = res_cur; + smpl.is_backend = cur_prefix; + backend_prefix = cur_prefix; - res = res && res_cur; + res = res && cur_prefix; } auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false); @@ -822,6 +835,23 @@ static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) { } } +static void llama_sampler_chain_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) { + const auto * src_chain = (const llama_sampler_chain *) src->ctx; + auto * dst_chain = (llama_sampler_chain *) dst->ctx; + + GGML_ASSERT(src_chain->samplers.size() == dst_chain->samplers.size()); + + for (size_t i = 0; i < src_chain->samplers.size(); ++i) { + llama_sampler_copy(src_chain->samplers[i].ptr, dst_chain->samplers[i].ptr); + } + + // note: is_init, n_nodes and is_backend belong to the current sampling graph + dst_chain->params = src_chain->params; + dst_chain->cur = src_chain->cur; + dst_chain->t_sample_us = src_chain->t_sample_us; + dst_chain->n_sample = src_chain->n_sample; +} + static struct llama_sampler_i llama_sampler_chain_i = { /* .name = */ llama_sampler_chain_name, /* .accept = */ llama_sampler_chain_accept, @@ -834,6 +864,7 @@ static struct llama_sampler_i llama_sampler_chain_i = { /* .backend_apply = */ llama_sampler_chain_backend_apply, /* .backend_set_input = */ llama_sampler_chain_backend_set_input, /* .backend_reset = */ llama_sampler_chain_backend_reset, + /* .copy_state = */ llama_sampler_chain_copy_state, }; struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) { @@ -1031,9 +1062,9 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_greedy *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1070,6 +1101,7 @@ static struct llama_sampler_i llama_sampler_greedy_i = { /* .backend_apply = */ llama_sampler_greedy_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_greedy() { @@ -1097,6 +1129,15 @@ struct llama_sampler_dist : public llama_sampler_backend { // inputs for the current sampling graph std::vector inp_uniforms; + + void copy_state(const llama_sampler_dist & src) { + // note: inp_uniforms and backend_transactional belong to the current sampling graph + seed_cur = src.seed_cur; + rng = src.rng; + rng_backend = src.rng_backend; + n_backend_draws_generated = src.n_backend_draws_generated; + n_backend_draws_committed = src.n_backend_draws_committed; + } }; static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) { @@ -1196,6 +1237,7 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample { auto * result_ctx = (llama_sampler_dist *) result->ctx; + result_ctx->seed_cur = ctx->seed_cur; result_ctx->rng = ctx->rng; result_ctx->backend_transactional = ctx->backend_transactional; result_ctx->rng_backend = ctx->rng_backend; @@ -1213,13 +1255,13 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_dist *) smpl->ctx; const bool res = llama_sampler_backend_support(smpl, buft); sctx->init(res); - sctx->backend_transactional = n_outputs_per_seq_max > 1; + sctx->backend_transactional = n_outputs_max_per_seq > 1; sctx->rng_backend = sctx->rng; sctx->n_backend_draws_generated = 0; sctx->n_backend_draws_committed = 0; @@ -1297,7 +1339,7 @@ static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) { 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 + // llama_sampler_dist which uses double precision (sampling from // std::uniform_real_distribution and // std::uniform_real_distribution with same rng will produce // different sequences). @@ -1349,6 +1391,7 @@ static struct llama_sampler_i llama_sampler_dist_i = { /* .backend_apply = */ llama_sampler_dist_backend_apply, /* .backend_set_input = */ llama_sampler_dist_backend_set_input, /* .backend_reset = */ llama_sampler_dist_backend_reset, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_dist(uint32_t seed) { @@ -1418,9 +1461,9 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_top_k *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1468,6 +1511,7 @@ static struct llama_sampler_i llama_sampler_top_k_i = { /* .backend_apply = */ llama_sampler_top_k_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_top_k(int32_t k) { @@ -1567,9 +1611,9 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_top_p *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1667,6 +1711,7 @@ static struct llama_sampler_i llama_sampler_top_p_i = { /* .backend_apply = */ llama_sampler_top_p_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) { @@ -1765,9 +1810,9 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_min_p *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1829,6 +1874,7 @@ static struct llama_sampler_i llama_sampler_min_p_i = { /* .backend_apply = */ llama_sampler_min_p_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) { @@ -1940,6 +1986,7 @@ static struct llama_sampler_i llama_sampler_typical_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) { @@ -2017,9 +2064,9 @@ static void llama_sampler_backend_temp_sampling( static bool llama_sampler_temp_backend_init( struct llama_sampler * smpl, ggml_backend_buffer_type_t buft, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_temp *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -2049,6 +2096,7 @@ static struct llama_sampler_i llama_sampler_temp_i = { /* .backend_apply = */ llama_sampler_temp_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_temp(float temp) { @@ -2163,9 +2211,9 @@ 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, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_temp_ext *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); const bool res = llama_sampler_backend_support(smpl, buft); @@ -2251,6 +2299,7 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = { /* .backend_apply = */ llama_sampler_temp_ext_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) { @@ -2359,6 +2408,7 @@ static struct llama_sampler_i llama_sampler_xtc_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) { @@ -2447,7 +2497,7 @@ static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sa // copy the state { - auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx; + auto * result_ctx = (llama_sampler_mirostat *) result->ctx; result_ctx->mu = ctx->mu; result_ctx->rng = ctx->rng; @@ -2479,6 +2529,7 @@ static struct llama_sampler_i llama_sampler_mirostat_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) { @@ -2584,6 +2635,7 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) { @@ -2706,6 +2758,7 @@ static struct llama_sampler_i llama_sampler_grammar_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; static struct llama_sampler * llama_sampler_init_grammar_impl( @@ -2821,6 +2874,12 @@ struct llama_sampler_penalties : public llama_sampler_backend { std::vector host_token_ids; std::vector host_counts; + void copy_state(const llama_sampler_penalties & src) { + // note: inp_token_ids/inp_counts belong to the current sampling graph + prev = src.prev; + token_count = src.token_count; + } + static bool is_disabled( int32_t penalty_last_n, float penalty_repeat, @@ -2951,10 +3010,10 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) { static bool llama_sampler_penalties_backend_init( struct llama_sampler * smpl, ggml_backend_buffer_type_t buft, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { auto * sctx = (llama_sampler_penalties *) smpl->ctx; - if (n_outputs_per_seq_max > 1) { + if (n_outputs_max_per_seq > 1) { sctx->init(false); return false; } @@ -3136,6 +3195,7 @@ static struct llama_sampler_i llama_sampler_penalties_i = { /* .backend_apply = */ llama_sampler_penalties_backend_apply, /* .backend_set_input = */ llama_sampler_penalties_backend_set_input, /* .backend_reset = */ llama_sampler_penalties_backend_reset, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_penalties( @@ -3232,6 +3292,7 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_top_n_sigma(float n) { @@ -3570,6 +3631,7 @@ static struct llama_sampler_i llama_sampler_dry_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_dry(const struct llama_vocab * vocab, float dry_multiplier, float dry_base, int32_t dry_allowed_length, int32_t dry_penalty_last_n, const char** seq_breakers, size_t num_breakers) { @@ -3790,6 +3852,7 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_adaptive_p( @@ -3945,9 +4008,9 @@ static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) static bool llama_sampler_logit_bias_backend_init( struct llama_sampler * smpl, ggml_backend_buffer_type_t buft, - uint32_t n_outputs_per_seq_max) { + uint32_t n_outputs_max_per_seq) { GGML_UNUSED(buft); - GGML_UNUSED(n_outputs_per_seq_max); + GGML_UNUSED(n_outputs_max_per_seq); auto * sctx = (llama_sampler_logit_bias *) smpl->ctx; @@ -3972,6 +4035,7 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = { /* .backend_apply = */ llama_sampler_logit_bias_backend_apply, /* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input, /* .backend_reset = */ llama_sampler_logit_bias_backend_reset, + /* .copy_state = */ llama_sampler_backend_copy_state, }; struct llama_sampler * llama_sampler_init_logit_bias( @@ -4216,6 +4280,7 @@ static struct llama_sampler_i llama_sampler_infill_i = { /* .backend_apply = */ nullptr, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) { @@ -4229,6 +4294,32 @@ struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * voca ); } +void llama_sampler_copy(const struct llama_sampler * src, struct llama_sampler * dst) { + if (!src || !dst || src == dst) { + return; + } + + GGML_ASSERT(src->iface == dst->iface && "llama_sampler_copy: cannot copy between different sampler types"); + + if (dst->iface->copy_state) { + dst->iface->copy_state(src, dst); + return; + } + + // build a temporary sampler carrying src's current state + llama_sampler * tmp = llama_sampler_clone(src); + + // free dst's old state (frees dst->ctx, including children for a chain) + if (dst->iface->free) { + dst->iface->free(dst); + } + + // transplant tmp's state into dst, then destroy the (now empty) temp shell + dst->ctx = tmp->ctx; + tmp->ctx = nullptr; + delete tmp; +} + // utils uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) { diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 9e1d13daf9..ba58f852eb 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -37,11 +37,11 @@ static void test(void) { { common_params base; base.n_parallel = 4; - base.n_sampling_outputs_per_seq_max = 8; + base.n_outputs_max_per_seq = 8; const auto draft = common_base_params_to_speculative(base); assert(draft.n_outputs_max == 4); - assert(draft.n_sampling_outputs_per_seq_max == 1); + assert(draft.n_outputs_max_per_seq == 1); } printf("test-arg-parser: make sure there is no duplicated arguments in any examples\n\n"); diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index f244c7e944..10bdd7ad8c 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -87,7 +87,7 @@ struct test_context { int32_t n_seq_max = -1, uint32_t n_outputs_max = 0, uint32_t n_ubatch = 0, - uint32_t n_sampling_outputs_per_seq_max = 1) { + uint32_t n_outputs_max_per_seq = 1) { auto * model = params.model.get(); GGML_ASSERT(model); @@ -100,7 +100,7 @@ struct test_context { cparams.n_ubatch = n_ubatch; } cparams.n_outputs_max = n_outputs_max; - cparams.n_sampling_outputs_per_seq_max = n_sampling_outputs_per_seq_max; + cparams.n_outputs_max_per_seq = n_outputs_max_per_seq; cparams.samplers = configs.data(); cparams.n_samplers = configs.size(); cparams.kv_unified = true; @@ -276,7 +276,7 @@ struct test_context { struct test_single_output_backend_sampler { bool backend_initialized = false; - uint32_t backend_outputs_per_seq_max = 0; + uint32_t backend_outputs_max_per_seq = 0; int backend_apply_count = 0; int apply_count = 0; }; @@ -296,10 +296,10 @@ static void test_single_output_backend_sampler_free(llama_sampler * smpl) { } static bool test_single_output_backend_sampler_backend_init( - llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_per_seq_max) { + llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_max_per_seq) { auto * ctx = (test_single_output_backend_sampler *) smpl->ctx; - ctx->backend_outputs_per_seq_max = n_outputs_per_seq_max; - if (n_outputs_per_seq_max > 1) { + ctx->backend_outputs_max_per_seq = n_outputs_max_per_seq; + if (n_outputs_max_per_seq > 1) { return false; } ctx->backend_initialized = true; @@ -324,6 +324,7 @@ static llama_sampler_i test_single_output_backend_sampler_i = { /* .backend_apply = */ test_single_output_backend_sampler_backend_apply, /* .backend_set_input = */ nullptr, /* .backend_reset = */ nullptr, + /* .copy_state = */ nullptr, }; static llama_sampler * test_single_output_backend_sampler_init( @@ -1771,8 +1772,7 @@ static void test_backend_multi_output_dist_transaction(const test_params & param verify_random(0, randoms[2]); llama_batch_free(batch); - GGML_ASSERT(llama_set_sampler(test_ctx.ctx.get(), seq_id, saved.get())); - chain = std::move(saved); + llama_sampler_copy(saved.get(), chain.get()); batch = decode(); verify_random(0, randoms[2]); @@ -1955,7 +1955,7 @@ static void test_backend_multi_output_cpu_suffix(const test_params & params) { GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); GGML_ASSERT(sampler_ctx->backend_initialized); - GGML_ASSERT(sampler_ctx->backend_outputs_per_seq_max == 1); + GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 1); 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); @@ -1976,7 +1976,7 @@ static void test_backend_multi_output_cpu_suffix(const test_params & params) { GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); GGML_ASSERT(!sampler_ctx->backend_initialized); - GGML_ASSERT(sampler_ctx->backend_outputs_per_seq_max == 2); + GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 2); 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); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a25c6602e2..3b5f6a1218 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -302,7 +302,6 @@ 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,7 +363,6 @@ struct server_slot { task.reset(); llama_set_sampler(ctx_tgt, id, nullptr); - backend_sampling = false; // clear alora start alora_invocation_start = -1; @@ -1066,7 +1064,7 @@ private: params_base = params; const auto output_limits = server_output_limits(params_base); params_base.n_outputs_max = output_limits.total; - params_base.n_sampling_outputs_per_seq_max = output_limits.per_seq; + params_base.n_outputs_max_per_seq = output_limits.per_seq; const bool has_mmproj = !params.mmproj.path.empty(); const bool has_draft = params.speculative.has_dft(); @@ -1842,10 +1840,9 @@ private: // TODO: tmp until backend sampling is fully implemented if (use_backend_sampling) { - slot.backend_sampling = llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get())); + 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()); @@ -3914,12 +3911,7 @@ private: slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); slot.prompt.tokens.keep_first(ckpt.n_tokens); - if (slot.backend_sampling) { - slot.backend_sampling = llama_set_sampler( - slot.ctx_tgt, slot.id, common_sampler_get(smpl_save.get())); - } - - slot.smpl = std::move(smpl_save); + common_sampler_copy(smpl_save.get(), slot.smpl.get()); return; }