diff --git a/common/llguidance.cpp b/common/llguidance.cpp index d58f147a76..51b3108d4e 100644 --- a/common/llguidance.cpp +++ b/common/llguidance.cpp @@ -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, }; diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 1fe242d062..ce31579939 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -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, }; diff --git a/common/speculative.cpp b/common/speculative.cpp index 70dc0ac3b1..b67c63a732 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2377,6 +2377,11 @@ common_speculative_init_result_ptr common_speculative_init_from_params(common_pa return std::make_unique(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(n_batch, n_outputs); +} + // initialization of the speculative decoding system // common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) { diff --git a/common/speculative.h b/common/speculative.h index 062bf20931..f7acc7d1f4 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -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); diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp index 2d4c0e528d..e685e15958 100644 --- a/examples/lookup/lookup.cpp +++ b/examples/lookup/lookup.cpp @@ -3,9 +3,11 @@ #include "common.h" #include "ngram-cache.h" #include "sampling.h" +#include "speculative.h" #include "log.h" #include "llama.h" +#include #include #include #include @@ -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); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index d87ba48beb..a85e262c68 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -5,6 +5,7 @@ #include "log.h" #include "llama.h" +#include #include #include #include @@ -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); diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index f7fa5e3060..6a9c9b1c53 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -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; diff --git a/include/llama.h b/include/llama.h index a14498925f..1abc5c3947 100644 --- a/include/llama.h +++ b/include/llama.h @@ -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); }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 19cca7df1e..b6ffdfe1ac 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -10,6 +10,7 @@ #include "llama-mmap.h" #include "llama-model.h" #include "llama-ext.h" +#include "llama-sampler.h" #include "llama.h" #include @@ -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 build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) { - std::map 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 & tensor_map, - const buffer_view & sampled, - const std::map & 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 & tensor_map, - const buffer_view & dst, +template +static void copy_tensor_async_rows( + const std::vector & tensors, + const buffer_view & dst, size_t stride, - std::vector & counts, - const std::map & seq_to_row, - ggml_backend_sched_t sched) { + uint32_t row_offset, + ggml_backend_sched_t sched, + std::vector * 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 & tensor_map, - const buffer_view & dst, - size_t stride, - std::vector & counts, - const std::map & 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 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(n_tokens * 40, 32u * model.n_tensors()); + res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else { + res = std::max(1024u, 8u*model.n_tensors()); + for (const auto & lora : model.loras) { + res += lora->get_n_nodes(); + } } - uint32_t res = std::max(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 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(); diff --git a/src/llama-cparams.h b/src/llama-cparams.h index 5018170ed8..6f0ae54710 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -52,6 +52,7 @@ struct llama_cparams { bool op_offload; bool kv_unified; bool pipeline_parallel; + bool sampler_backend_require_multi_output; std::vector embeddings_layer_inp; // [n_layer()] extract input embeddings for layer diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 2be3b75fb9..55d8580246 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -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(samplers); res->add_input(std::move(inp_sampling)); - std::map seq_to_logit_row; - int32_t logit_row_idx = 0; - - for (uint32_t i = 0; i < ubatch.n_tokens; i++) { + std::map> 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 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); } } } diff --git a/src/llama-graph.h b/src/llama-graph.h index 32d8d395aa..75bc0fe80d 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -904,10 +904,10 @@ public: std::vector t_layer_inp; - std::map t_sampled_logits; - std::map t_candidates; - std::map t_sampled; - std::map t_sampled_probs; + std::vector t_sampled; + std::vector t_sampled_probs; + std::vector t_sampled_logits; + std::vector t_candidates; std::vector inputs; std::vector fused_nodes; diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index e550fbe4ae..0bb7484b1d 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -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(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 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 with same rng will produce // different sequences). std::uniform_real_distribution 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) { diff --git a/src/llama-sampler.h b/src/llama-sampler.h index 9292075146..9cd84dccb4 100644 --- a/src/llama-sampler.h +++ b/src/llama-sampler.h @@ -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, diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 50db297274..1a72e635e0 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -2,7 +2,9 @@ #include "common.h" #include "download.h" #include "llama.h" +#include "speculative.h" +#include #include #include #include @@ -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::max(), + std::numeric_limits::max(), + std::numeric_limits::max()) == std::numeric_limits::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 { diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 1165f46f0c..763be68f37 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -80,7 +81,12 @@ struct test_context { std::unordered_map seq_positions; std::unordered_map last_batch_info; - test_context(const test_params & params, std::vector & configs, int32_t n_seq_max = -1) { + test_context( + const test_params & params, + std::vector & 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 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 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 tokens; - tokens.push_back(llama_vocab_bos(test_ctx.vocab)); - - std::vector 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 configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 4); + + std::vector 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(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 configs = {{ seq_id, chain.get() }}; + test_context test_ctx(params, configs, 1, 4, 2); + + std::vector reference_configs; + test_context reference_ctx(params, reference_configs, 1, 4, 2); + llama_sampler_ptr reference_filters = make_filter_chain(); + + std::vector reference_data(n_vocab); + std::mt19937 reference_rng(seed); + std::uniform_real_distribution 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 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 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 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 }, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 38d2e5c7a0..ea9664fb57 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -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(1, std::min(n_batch, n_outputs)); + return std::max(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; } diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index c6568479ca..5fc338ad78 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -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",