diff --git a/common/common.cpp b/common/common.cpp index ffe3e7761b..84d4cdaa86 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1639,6 +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_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 2e15ec3f81..03255108c0 100644 --- a/common/common.h +++ b/common/common.h @@ -447,6 +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 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/speculative.cpp b/common/speculative.cpp index b67c63a732..1e8a453423 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2382,6 +2382,11 @@ int32_t common_speculative_n_outputs_max(int32_t n_batch, int32_t n_parallel, in return std::min(n_batch, n_outputs); } +int32_t common_speculative_n_outputs_per_seq_max(int32_t n_batch, int32_t n_draft) { + const int64_t n_outputs = 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 f7acc7d1f4..d710d356dc 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -28,6 +28,9 @@ 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); +// return the max number of outputs per sequence needed for speculative decoding +int32_t common_speculative_n_outputs_per_seq_max(int32_t n_batch, 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 e685e15958..bdf4db70d8 100644 --- a/examples/lookup/lookup.cpp +++ b/examples/lookup/lookup.cpp @@ -30,6 +30,7 @@ int main(int argc, char ** argv){ 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); + params.n_sampling_outputs_per_seq_max = common_speculative_n_outputs_per_seq_max(params.n_batch, n_draft); // init llama.cpp llama_backend_init(); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index a85e262c68..25faa43bec 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -32,6 +32,8 @@ int main(int argc, char ** argv) { params.n_outputs_max = common_speculative_n_outputs_max( params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative)); + params.n_sampling_outputs_per_seq_max = common_speculative_n_outputs_per_seq_max( + params.n_batch, common_speculative_n_max(¶ms.speculative)); // init llama.cpp llama_backend_init(); @@ -59,6 +61,9 @@ 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.devices = params_spec.devices; params_dft.model = params_spec.mparams; params_dft.n_gpu_layers = params_spec.n_gpu_layers; diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index 6a9c9b1c53..f2033c0458 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -60,6 +60,8 @@ int main(int argc, char ** argv) { params.n_outputs_max = common_speculative_n_outputs_max( params.n_batch, params.n_parallel, params.speculative.draft.n_max); + params.n_sampling_outputs_per_seq_max = common_speculative_n_outputs_per_seq_max( + params.n_batch, 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; @@ -87,6 +89,8 @@ int main(int argc, char ** argv) { params.devices = params.speculative.draft.devices; 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; 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 1abc5c3947..75dc50652e 100644 --- a/include/llama.h +++ b/include/llama.h @@ -354,6 +354,7 @@ extern "C" { 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 @@ -1275,7 +1276,7 @@ extern "C" { bool (*backend_init)( struct llama_sampler * smpl, ggml_backend_buffer_type_t buft, - bool require_multi_output); + uint32_t n_outputs_per_seq_max); // call after .backend_apply() void (*backend_accept)( diff --git a/src/llama-context.cpp b/src/llama-context.cpp index b6ffdfe1ac..337484944b 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -100,7 +100,6 @@ 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)) { @@ -161,25 +160,6 @@ llama_context::llama_context( } } - // 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 - // re-reserve when graph nodes change. - if (params.samplers != nullptr && params.n_samplers > 0) { - for (size_t i = 0; i < params.n_samplers; ++i) { - const auto & config = params.samplers[i]; - - if (llama_sampler_chain_get(config.sampler, -1) == nullptr) { - throw std::runtime_error("the backend samplers must be of type llama_sampler_chain"); - } - - if (set_sampler(config.seq_id, config.sampler)) { - const int n_samplers = llama_sampler_chain_n(config.sampler); - - LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers); - } - } - } - auto rope_scaling_type = params.rope_scaling_type; if (rope_scaling_type == LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED) { rope_scaling_type = hparams.rope_scaling_type_train; @@ -267,6 +247,27 @@ 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); + + // 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 + // re-reserve when graph nodes change. + if (params.samplers != nullptr && params.n_samplers > 0) { + for (size_t i = 0; i < params.n_samplers; ++i) { + const auto & config = params.samplers[i]; + + if (llama_sampler_chain_get(config.sampler, -1) == nullptr) { + throw std::runtime_error("the backend samplers must be of type llama_sampler_chain"); + } + + if (set_sampler(config.seq_id, config.sampler)) { + const int n_samplers = llama_sampler_chain_n(config.sampler); + + LLAMA_LOG_INFO("%s: setting backend sampler for seq_id %d (n = %d)\n", __func__, config.seq_id, n_samplers); + } + } + } cparams.op_offload = params.op_offload; cparams.kv_unified = params.kv_unified; @@ -314,6 +315,7 @@ llama_context::llama_context( 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); 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", @@ -1233,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.sampler_backend_require_multi_output); + sampler->iface->backend_init(sampler, buft, cparams.n_sampling_outputs_per_seq_max); sampling.samplers[seq_id] = sampler; @@ -1674,11 +1676,11 @@ int llama_context::decode(const llama_batch & batch_inp) { seq_output_count[seq_id]++; 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]); + if (sampler != sampling.samplers.end() && + seq_output_count[seq_id] > (int32_t) cparams.n_sampling_outputs_per_seq_max) { + 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, seq_output_count[seq_id]); return -1; } } @@ -2303,7 +2305,7 @@ 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.sampler_backend_require_multi_output) { + if (cparams.n_sampling_outputs_per_seq_max > 1) { n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes); } } @@ -2359,7 +2361,7 @@ ggml_cgraph * llama_context::graph_reserve( 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)) { + cparams.n_sampling_outputs_per_seq_max == 1)) { continue; } @@ -3453,6 +3455,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_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 6f0ae54710..dbd4021c9f 100644 --- a/src/llama-cparams.h +++ b/src/llama-cparams.h @@ -15,6 +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; int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing @@ -52,7 +53,6 @@ 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-sampler.cpp b/src/llama-sampler.cpp index 9175351754..6e96ff70c7 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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { GGML_UNUSED(smpl); GGML_UNUSED(buft); - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); return true; } @@ -715,7 +715,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * chain = (llama_sampler_chain *) smpl->ctx; GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice"); @@ -731,9 +731,9 @@ static bool llama_sampler_chain_backend_init( // 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 output mode + // - support the requested per-sequence output limit if (res_cur && smpl.ptr->iface->backend_init) { - if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, require_multi_output)) { + if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_per_seq_max)) { res_cur = false; } } else { @@ -1025,9 +1025,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_greedy *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1191,9 +1191,9 @@ 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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_dist *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1350,9 +1350,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_top_k *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1499,9 +1499,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_top_p *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1697,9 +1697,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_min_p *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1949,9 +1949,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_temp *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -2095,9 +2095,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_temp_ext *) smpl->ctx; - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -3865,9 +3865,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, - bool require_multi_output) { + uint32_t n_outputs_per_seq_max) { GGML_UNUSED(buft); - GGML_UNUSED(require_multi_output); + GGML_UNUSED(n_outputs_per_seq_max); auto * sctx = (llama_sampler_logit_bias *) smpl->ctx; diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 1a72e635e0..71a2fec5ca 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -23,6 +23,12 @@ static void test(void) { std::numeric_limits::max(), std::numeric_limits::max(), std::numeric_limits::max()) == std::numeric_limits::max()); + assert(common_speculative_n_outputs_per_seq_max(16, 3) == 4); + assert(common_speculative_n_outputs_per_seq_max(16, -1) == 1); + assert(common_speculative_n_outputs_per_seq_max(2, 3) == 2); + assert(common_speculative_n_outputs_per_seq_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++) { diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 763be68f37..1891442359 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -86,7 +86,8 @@ struct test_context { std::vector & configs, int32_t n_seq_max = -1, uint32_t n_outputs_max = 0, - uint32_t n_ubatch = 0) { + uint32_t n_ubatch = 0, + uint32_t n_sampling_outputs_per_seq_max = 1) { auto * model = params.model.get(); GGML_ASSERT(model); @@ -99,6 +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.samplers = configs.data(); cparams.n_samplers = configs.size(); cparams.kv_unified = true; @@ -274,7 +276,7 @@ struct test_context { struct test_single_output_backend_sampler { bool backend_initialized = false; - bool backend_require_multi_output = false; + uint32_t backend_outputs_per_seq_max = 0; int backend_apply_count = 0; int apply_count = 0; }; @@ -294,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*/, bool require_multi_output) { + llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_per_seq_max) { auto * ctx = (test_single_output_backend_sampler *) smpl->ctx; - ctx->backend_require_multi_output = require_multi_output; - if (require_multi_output) { + ctx->backend_outputs_per_seq_max = n_outputs_per_seq_max; + if (n_outputs_per_seq_max > 1) { return false; } ctx->backend_initialized = true; @@ -1596,26 +1598,27 @@ static void test_backend_cpu_mixed_batch(const test_params & params) { printf("backend-cpu mixed batch test PASSED\n"); } -static void test_backend_multi_output_disabled(const test_params & params) { +static void test_backend_multi_output_limit(const test_params & params) { const llama_seq_id seq_id = 0; 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, configs, 1, 3, 0, 2); - 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(3, 0, 1); + for (int i = 0; i < 3; ++i) { + common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), i, { seq_id }, true); + } - printf(">>> test_backend_multi_output_disabled expected error start:\n"); + printf(">>> test_backend_multi_output_limit expected error start:\n"); const int ret = llama_decode(test_ctx.ctx.get(), batch); - GGML_ASSERT(ret != 0 && "llama_decode should reject multiple outputs for one sequence"); - printf("<<< test_backend_multi_output_disabled expected error end.\n"); + GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit"); + printf("<<< test_backend_multi_output_limit expected error end.\n"); llama_batch_free(batch); - printf("backend multi-output disabled test PASSED\n"); + printf("backend multi-output limit test PASSED\n"); } // greedy is a stateless terminal selector; verify multi-output backend argmax @@ -1628,7 +1631,7 @@ static void test_backend_multi_output_greedy(const test_params & params) { 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); + test_context test_ctx(params, configs, 4, 4, 0, 4); std::vector reference_configs; test_context reference_ctx(params, reference_configs, 1, 4); @@ -1686,7 +1689,7 @@ static void test_backend_multi_output_sampling_chain(const test_params & params) 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); + test_context test_ctx(params, configs, 1, 4, 2, 4); std::vector reference_configs; test_context reference_ctx(params, reference_configs, 1, 4, 2); @@ -1814,14 +1817,14 @@ static void test_backend_multi_output_cpu_suffix(const test_params & params) { 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); + test_context test_ctx(params, configs, 1, 1, 0, 4); 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_outputs_per_seq_max == 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); @@ -1836,7 +1839,7 @@ static void test_backend_multi_output_cpu_suffix(const test_params & params) { 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); + test_context test_ctx(params, configs, 1, 2, 0, 0); llama_batch batch = llama_batch_init(2, 0, 1); for (int i = 0; i < 2; ++i) { @@ -1845,7 +1848,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_require_multi_output); + GGML_ASSERT(sampler_ctx->backend_outputs_per_seq_max == 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); @@ -1879,7 +1882,7 @@ 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 }, - { "multi_output_disabled", test_backend_multi_output_disabled, true }, + { "multi_output_limit", test_backend_multi_output_limit, 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 }, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index ea9664fb57..5c330d101a 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -51,6 +51,18 @@ static uint32_t server_n_outputs_max(const common_params & params) { return std::max(1, n_outputs); } +static uint32_t server_n_sampling_outputs_per_seq_max(const common_params & params) { + if (params.embedding || + (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { + return 1; + } + + const int32_t n_outputs = common_speculative_n_outputs_per_seq_max( + params.n_batch, common_speculative_n_max(¶ms.speculative)); + + return std::max(1, n_outputs); +} + // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 enum slot_state { SLOT_STATE_IDLE, @@ -1063,6 +1075,7 @@ private: params_base = params; params_base.n_outputs_max = server_n_outputs_max(params_base); + params_base.n_sampling_outputs_per_seq_max = server_n_sampling_outputs_per_seq_max(params_base); const bool has_mmproj = !params.mmproj.path.empty(); const bool has_draft = params.speculative.has_dft(); @@ -1143,6 +1156,7 @@ private: bool measure_model_bytes = has_draft; common_params params_dft = common_base_params_to_speculative(params_base); + params_dft.n_sampling_outputs_per_seq_max = 1; auto mparams_dft = common_model_params_to_llama(params_dft); auto cparams_dft = common_context_params_to_llama(params_dft); @@ -1231,6 +1245,7 @@ private: // progress callback params_dft.load_progress_callback = load_progress_callback; params_dft.load_progress_callback_user_data = &load_progress_spec; + params_dft.n_sampling_outputs_per_seq_max = 1; spec_init = common_speculative_init_from_params(params_dft, model_tgt, ctx_tgt); model_dft = spec_init->model();