diff --git a/common/llguidance.cpp b/common/llguidance.cpp index 51b3108d4e..f72dbd1c1e 100644 --- a/common/llguidance.cpp +++ b/common/llguidance.cpp @@ -115,8 +115,8 @@ static llama_sampler_i llama_sampler_llg_i = { /* .backend_init = */ NULL, /* .backend_accept = */ NULL, /* .backend_apply = */ NULL, - /* .backend_reset = */ NULL, /* .backend_set_input = */ NULL, + /* .backend_reset = */ 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 ce31579939..58ad2bc07f 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -216,8 +216,8 @@ 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, + /* .backend_reset = */ nullptr, }; static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) { diff --git a/common/speculative.cpp b/common/speculative.cpp index 1e8a453423..a074ab608b 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -2292,6 +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; return result; } @@ -2377,14 +2378,15 @@ 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); -} +common_speculative_output_limits common_speculative_get_output_limits( + int32_t n_batch, int32_t n_parallel, int32_t n_draft) { + const int64_t per_seq = 1 + (int64_t) std::max(0, n_draft); + const int64_t total = (int64_t) n_parallel * per_seq; -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); + return { + /* .total = */ (int32_t) std::min(n_batch, total), + /* .per_seq = */ (int32_t) std::min(n_batch, per_seq), + }; } // initialization of the speculative decoding system diff --git a/common/speculative.h b/common/speculative.h index d710d356dc..c6986affda 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -25,11 +25,14 @@ 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); +struct common_speculative_output_limits { + int32_t total; + int32_t per_seq; +}; -// 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); +// return the output limits needed for speculative decoding +common_speculative_output_limits common_speculative_get_output_limits( + 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); diff --git a/examples/lookup/lookup.cpp b/examples/lookup/lookup.cpp index bdf4db70d8..75ee27c299 100644 --- a/examples/lookup/lookup.cpp +++ b/examples/lookup/lookup.cpp @@ -29,8 +29,9 @@ 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); - params.n_sampling_outputs_per_seq_max = common_speculative_n_outputs_per_seq_max(params.n_batch, n_draft); + 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; // init llama.cpp llama_backend_init(); diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp index 25faa43bec..c30fe11531 100644 --- a/examples/speculative-simple/speculative-simple.cpp +++ b/examples/speculative-simple/speculative-simple.cpp @@ -30,10 +30,10 @@ int main(int argc, char ** argv) { return 1; } - params.n_outputs_max = common_speculative_n_outputs_max( + const auto output_limits = common_speculative_get_output_limits( 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)); + params.n_outputs_max = output_limits.total; + params.n_sampling_outputs_per_seq_max = output_limits.per_seq; // init llama.cpp llama_backend_init(); diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index f2033c0458..b6838b0675 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -58,10 +58,10 @@ 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( + const auto output_limits = common_speculative_get_output_limits( 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); + params.n_outputs_max = output_limits.total; + params.n_sampling_outputs_per_seq_max = 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; diff --git a/include/llama.h b/include/llama.h index 35cab4d39d..5e7880dd8d 100644 --- a/include/llama.h +++ b/include/llama.h @@ -1057,6 +1057,7 @@ extern "C" { // Get the backend sampled token for the ith token. // With multiple outputs, sampler state advances when the token is accepted, // not when it is read through this function. + // When accepting multiple outputs, accept a contiguous prefix in output order. // Returns LLAMA_TOKEN_NULL if no token was sampled. LLAMA_API llama_token llama_get_sampled_token_ith(struct llama_context * ctx, int32_t i); @@ -1294,11 +1295,11 @@ 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); + + // called before rebuilding a sampling graph to clear graph-owned tensor references + void (*backend_reset)(struct llama_sampler * smpl); }; struct llama_sampler { @@ -1508,6 +1509,7 @@ extern "C" { LLAMA_API uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl); /// @details Sample and accept a token from the idx-th output of the last evaluation + // For multiple outputs from one sampler, call this function in output order without gaps. // // Shorthand for: // const auto * logits = llama_get_logits_ith(ctx, idx); diff --git a/src/llama-context.cpp b/src/llama-context.cpp index bda3554d65..61f4df06db 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2386,27 +2386,16 @@ 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); - // activate each configured sampler once - if (n_sampling_outputs_per_seq > 0) { - for (uint32_t s : sampler_seqs) { - if (n_outputs_set >= n_outputs) { - break; - } - - ubatch.output[s * ubatch.n_seq_tokens] = true; - ++n_outputs_set; - } - } - - // add the remaining valid sampling rows - for (uint32_t t = 1; t < n_sampling_outputs_per_seq && n_outputs_set < n_outputs; ++t) { - for (uint32_t s : sampler_seqs) { - if (n_outputs_set >= n_outputs) { - break; - } + // select sampling rows in round-robin order across sampler sequences + if (!sampler_seqs.empty()) { + const uint32_t n_sampler_seqs = sampler_seqs.size(); + n_outputs_set = std::min( + n_outputs, (uint64_t) n_sampler_seqs * n_sampling_outputs_per_seq); + for (uint32_t i = 0; i < n_outputs_set; ++i) { + const uint32_t s = sampler_seqs[i % n_sampler_seqs]; + const uint32_t t = i / n_sampler_seqs; ubatch.output[s * ubatch.n_seq_tokens + t] = true; - ++n_outputs_set; } } diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 6cf5971683..b3f57f123c 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -512,8 +512,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_empty(const char * name) { @@ -831,8 +831,8 @@ 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, + /* .backend_reset = */ llama_sampler_chain_backend_reset, }; struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) { @@ -1067,8 +1067,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_greedy() { @@ -1088,11 +1088,11 @@ struct llama_sampler_dist : public llama_sampler_backend { std::mt19937 rng; - // multi-output backend draws are committed when their tokens are accepted + // multi-output backend draws are committed as an accepted prefix bool backend_transactional; std::mt19937 rng_backend; - size_t n_backend_generated; - size_t n_backend_accepted; + size_t n_backend_draws_generated; + size_t n_backend_draws_committed; // inputs for the current sampling graph std::vector inp_uniforms; @@ -1183,8 +1183,8 @@ static void llama_sampler_dist_reset(struct llama_sampler * smpl) { ctx->seed_cur = get_rng_seed(ctx->seed); ctx->rng.seed(ctx->seed_cur); ctx->rng_backend = ctx->rng; - ctx->n_backend_generated = 0; - ctx->n_backend_accepted = 0; + ctx->n_backend_draws_generated = 0; + ctx->n_backend_draws_committed = 0; } static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sampler * smpl) { @@ -1195,11 +1195,11 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample { auto * result_ctx = (llama_sampler_dist *) result->ctx; - result_ctx->rng = ctx->rng; - result_ctx->backend_transactional = ctx->backend_transactional; - result_ctx->rng_backend = ctx->rng_backend; - result_ctx->n_backend_generated = ctx->n_backend_generated; - result_ctx->n_backend_accepted = ctx->n_backend_accepted; + result_ctx->rng = ctx->rng; + result_ctx->backend_transactional = ctx->backend_transactional; + result_ctx->rng_backend = ctx->rng_backend; + result_ctx->n_backend_draws_generated = ctx->n_backend_draws_generated; + result_ctx->n_backend_draws_committed = ctx->n_backend_draws_committed; } return result; @@ -1220,8 +1220,8 @@ static bool llama_sampler_dist_backend_init( sctx->init(res); sctx->backend_transactional = n_outputs_per_seq_max > 1; sctx->rng_backend = sctx->rng; - sctx->n_backend_generated = 0; - sctx->n_backend_accepted = 0; + sctx->n_backend_draws_generated = 0; + sctx->n_backend_draws_committed = 0; return res; } @@ -1311,7 +1311,7 @@ static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) { ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float)); if (sctx->backend_transactional) { - ++sctx->n_backend_generated; + ++sctx->n_backend_draws_generated; } } } @@ -1326,13 +1326,14 @@ static void llama_sampler_dist_accept(struct llama_sampler * smpl, llama_token t auto * sctx = (llama_sampler_dist *) smpl->ctx; - if (!sctx->backend_transactional || sctx->n_backend_accepted >= sctx->n_backend_generated) { + if (!sctx->backend_transactional || + sctx->n_backend_draws_committed >= sctx->n_backend_draws_generated) { return; } std::uniform_real_distribution dist(0.0f, 1.0f); dist(sctx->rng); - ++sctx->n_backend_accepted; + ++sctx->n_backend_draws_committed; } static struct llama_sampler_i llama_sampler_dist_i = { @@ -1345,8 +1346,8 @@ 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, + /* .backend_reset = */ llama_sampler_dist_backend_reset, }; struct llama_sampler * llama_sampler_init_dist(uint32_t seed) { @@ -1355,14 +1356,14 @@ 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), - /* .backend_transactional = */ false, - /* .rng_backend = */ std::mt19937(seed_cur), - /* .n_backend_generated = */ 0, - /* .n_backend_accepted = */ 0, - /* .inp_uniforms = */ {}, + /* .seed = */ seed, + /* .seed_cur = */ seed_cur, + /* .rng = */ std::mt19937(seed_cur), + /* .backend_transactional = */ false, + /* .rng_backend = */ std::mt19937(seed_cur), + /* .n_backend_draws_generated = */ 0, + /* .n_backend_draws_committed = */ 0, + /* .inp_uniforms = */ {}, } ); } @@ -1382,8 +1383,8 @@ void llama_sampler_backend_begin(llama_sampler * sampler) { auto * ctx = (llama_sampler_dist *) sampler->ctx; if (ctx->backend_transactional) { ctx->rng_backend = ctx->rng; - ctx->n_backend_generated = 0; - ctx->n_backend_accepted = 0; + ctx->n_backend_draws_generated = 0; + ctx->n_backend_draws_committed = 0; } } } @@ -1464,8 +1465,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_top_k(int32_t k) { @@ -1663,8 +1664,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) { @@ -1825,8 +1826,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) { @@ -1936,8 +1937,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) { @@ -2045,8 +2046,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_temp(float temp) { @@ -2247,8 +2248,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) { @@ -2355,8 +2356,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) { @@ -2475,8 +2476,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) { @@ -2580,8 +2581,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) { @@ -2702,8 +2703,8 @@ 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, + /* .backend_reset = */ nullptr, }; static struct llama_sampler * llama_sampler_init_grammar_impl( @@ -3120,8 +3121,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_penalties( @@ -3216,8 +3217,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_top_n_sigma(float n) { @@ -3554,8 +3555,8 @@ 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, + /* .backend_reset = */ 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) { @@ -3774,8 +3775,8 @@ 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, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_adaptive_p( @@ -3956,8 +3957,8 @@ 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, + /* .backend_reset = */ llama_sampler_logit_bias_backend_reset, }; struct llama_sampler * llama_sampler_init_logit_bias( @@ -4200,8 +4201,8 @@ static struct llama_sampler_i llama_sampler_infill_i = { /* .backend_init = */ nullptr, /* .backend_accept = */ nullptr, /* .backend_apply = */ nullptr, - /* .backend_reset = */ nullptr, /* .backend_set_input = */ nullptr, + /* .backend_reset = */ nullptr, }; struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) { diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index 71a2fec5ca..69f4dfb65f 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -16,19 +16,33 @@ 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( + auto assert_output_limits = [](int32_t n_batch, int32_t n_parallel, int32_t n_draft, + int32_t total, int32_t per_seq) { + const auto limits = common_speculative_get_output_limits(n_batch, n_parallel, n_draft); + assert(limits.total == total); + assert(limits.per_seq == per_seq); + }; + + assert_output_limits(16, 2, 3, 8, 4); + assert_output_limits(16, 2, -1, 2, 1); + assert_output_limits( 4, 2, 3, 4, 4); + assert_output_limits( 2, 1, 3, 2, 2); + assert_output_limits( 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()); + std::numeric_limits::max(), + std::numeric_limits::max()); + + { + common_params base; + base.n_parallel = 4; + base.n_sampling_outputs_per_seq_max = 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); + } 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 da3100f263..a76cf192ba 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -322,8 +322,8 @@ static llama_sampler_i test_single_output_backend_sampler_i = { /* .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, + /* .backend_reset = */ nullptr, }; static llama_sampler * test_single_output_backend_sampler_init( diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index e22696ea78..a25c6602e2 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -39,28 +39,18 @@ using json = nlohmann::ordered_json; constexpr int HTTP_POLLING_SECONDS = 1; -static uint32_t server_n_outputs_max(const common_params & params) { +static common_speculative_output_limits server_output_limits(const common_params & params) { if (params.embedding || (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) { - return params.n_batch; + return { params.n_batch, 1 }; } - const int32_t n_outputs = common_speculative_n_outputs_max( + auto result = common_speculative_get_output_limits( params.n_batch, params.n_parallel, common_speculative_n_max(¶ms.speculative)); - 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); + result.total = std::max(1, result.total); + result.per_seq = std::max(1, result.per_seq); + return result; } // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 @@ -1074,8 +1064,9 @@ private: const bool is_resume = sleeping; 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 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; const bool has_mmproj = !params.mmproj.path.empty(); const bool has_draft = params.speculative.has_dft(); @@ -1156,7 +1147,6 @@ 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); @@ -1245,7 +1235,6 @@ 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();