diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 337484944b..ed4bf6a375 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1674,6 +1674,10 @@ int llama_context::decode(const llama_batch & batch_inp) { for (int32_t s = 0; s < ns; ++s) { const llama_seq_id seq_id = batch_inp.seq_id ? batch_inp.seq_id[i][s] : 0; + if (seq_id < 0 || (uint32_t) seq_id >= n_seq_max) { + continue; + } + seq_output_count[seq_id]++; auto sampler = sampling.samplers.find(seq_id); if (sampler != sampling.samplers.end() && @@ -2310,9 +2314,13 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { } } + const uint32_t n_sampling_outputs_max = std::min( + std::min(n_tokens, cparams.n_outputs_max), + (uint64_t) cparams.n_seq_max * cparams.n_sampling_outputs_per_seq_max); + res += n_sampling_nodes; - if (cparams.n_outputs_max > 1) { - res += (cparams.n_outputs_max - 1) * n_sampling_nodes_max; + if (n_sampling_outputs_max > 1) { + res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max; } return res; } @@ -2361,7 +2369,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.n_sampling_outputs_per_seq_max == 1)) { + t >= cparams.n_sampling_outputs_per_seq_max)) { continue; } diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 6e96ff70c7..3f82bd768d 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -603,6 +603,10 @@ static llama_sampler_backend_probe llama_sampler_backend_probe_graph( } } + if (sampler->iface->backend_reset) { + sampler->iface->backend_reset(sampler); + } + return { std::move(ctx_ptr), gf }; } @@ -1085,6 +1089,8 @@ struct llama_sampler_dist : public llama_sampler_backend { // inputs for the current sampling graph std::vector inp_uniforms; + + bool copy_candidates = false; }; static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) { @@ -1193,7 +1199,8 @@ static bool llama_sampler_dist_backend_init( ggml_backend_buffer_type_t buft, uint32_t n_outputs_per_seq_max) { auto * sctx = (llama_sampler_dist *) smpl->ctx; - GGML_UNUSED(n_outputs_per_seq_max); + + sctx->copy_candidates = n_outputs_per_seq_max > 1; const bool res = llama_sampler_backend_support(smpl, buft); @@ -1261,9 +1268,11 @@ 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"); + if (sctx->copy_candidates) { + // 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; diff --git a/tests/test-backend-sampler.cpp b/tests/test-backend-sampler.cpp index 1891442359..5a10156499 100644 --- a/tests/test-backend-sampler.cpp +++ b/tests/test-backend-sampler.cpp @@ -1666,6 +1666,99 @@ static void test_backend_multi_output_greedy(const test_params & params) { printf("backend multi-output greedy test PASSED\n"); } +static void test_backend_multi_sequence_multi_output_dist(const test_params & params) { + const llama_vocab * vocab = llama_model_get_vocab(params.model.get()); + const int32_t n_vocab = llama_vocab_n_tokens(vocab); + const uint32_t seeds[] = { 88, 1337 }; + // reduce the chance that swapped random inputs select the same token + const float temp = 10.0f; + + llama_sampler_ptr chain_0(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_ptr chain_1(llama_sampler_chain_init(llama_sampler_chain_default_params())); + llama_sampler_chain_add(chain_0.get(), llama_sampler_init_temp(temp)); + llama_sampler_chain_add(chain_0.get(), llama_sampler_init_dist(seeds[0])); + llama_sampler_chain_add(chain_1.get(), llama_sampler_init_temp(temp)); + llama_sampler_chain_add(chain_1.get(), llama_sampler_init_dist(seeds[1])); + std::vector configs = { + { 0, chain_0.get() }, + { 1, chain_1.get() }, + }; + test_context test_ctx(params, configs, 2, 6, 0, 3); + + std::vector reference_configs; + test_context reference_ctx(params, reference_configs, 2, 6); + + const llama_token seq_tokens[2][3] = { + { llama_vocab_bos(vocab), llama_vocab_eos(vocab), llama_vocab_bos(vocab) }, + { llama_vocab_eos(vocab), llama_vocab_bos(vocab), llama_vocab_eos(vocab) }, + }; + + llama_batch batch = llama_batch_init(6, 0, 1); + for (int pos = 0; pos < 3; ++pos) { + common_batch_add(batch, seq_tokens[0][pos], pos, { 0 }, true); + common_batch_add(batch, seq_tokens[1][pos], pos, { 1 }, true); + } + + GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); + GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0); + + std::mt19937 reference_rngs[] = { + std::mt19937(seeds[0]), + std::mt19937(seeds[1]), + }; + std::uniform_real_distribution reference_dist(0.0, 1.0); + int outputs_per_seq[] = { 0, 0 }; + + for (int i = 0; i < batch.n_tokens; ++i) { + const llama_seq_id seq_id = batch.seq_id[i][0]; + GGML_ASSERT(seq_id == 0 || seq_id == 1); + outputs_per_seq[seq_id]++; + + 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 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 float * reference_logits = llama_get_logits_ith(reference_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(reference_logits != nullptr); + GGML_ASSERT(n_logits == (uint32_t) n_vocab); + GGML_ASSERT(n_probs == (uint32_t) n_vocab); + + float prob_sum = 0.0f; + float cumsum_before = 0.0f; + for (llama_token token = 0; token < n_vocab; ++token) { + const float expected_logit = reference_logits[token] / temp; + const float tolerance = 1e-4f * std::max(1.0f, std::fabs(expected_logit)); + GGML_ASSERT(std::fabs(sampled_logits[token] - expected_logit) <= tolerance); + GGML_ASSERT(std::isfinite(sampled_probs[token])); + GGML_ASSERT(sampled_probs[token] >= 0.0f); + + prob_sum += sampled_probs[token]; + if (token < backend_token) { + cumsum_before += sampled_probs[token]; + } + } + + GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f); + + const float rnd = reference_dist(reference_rngs[seq_id]); + const float cumsum_sampled = cumsum_before + sampled_probs[backend_token]; + GGML_ASSERT(rnd >= cumsum_before - 1e-4f); + GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f); + } + + GGML_ASSERT(outputs_per_seq[0] == 3); + GGML_ASSERT(outputs_per_seq[1] == 3); + + llama_batch_free(batch); + + printf("backend multi-sequence multi-output dist 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; @@ -1698,12 +1791,15 @@ static void test_backend_multi_output_sampling_chain(const test_params & params) 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; + const int n_outputs_per_round[] = { 4, 3, 4, 4 }; + int32_t n_reused_before_repeat = -1; + int32_t pos = 0; - 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); + for (int round = 0; round < (int) (sizeof(n_outputs_per_round) / sizeof(n_outputs_per_round[0])); ++round) { + const int n_outputs = n_outputs_per_round[round]; + llama_batch batch = llama_batch_init(n_outputs, 0, 1); + for (int i = 0; i < n_outputs; ++i) { + common_batch_add(batch, llama_vocab_bos(vocab), pos++, { seq_id }, true); } GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0); @@ -1795,10 +1891,10 @@ static void test_backend_multi_output_sampling_chain(const test_params & params) 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); + if (round == 2) { + n_reused_before_repeat = n_reused; + } else if (round == 3) { + GGML_ASSERT(n_reused > n_reused_before_repeat); } } @@ -1884,6 +1980,7 @@ static const backend_test_case BACKEND_TESTS[] = { { "set_sampler", test_backend_set_sampler, true }, { "multi_output_limit", test_backend_multi_output_limit, true }, { "multi_output_greedy", test_backend_multi_output_greedy, true }, + { "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, 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 }, diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 5c330d101a..a9c6e638b1 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -222,6 +222,7 @@ struct server_slot { std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + common_sampler_ptr spec_smpl_save; // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 @@ -359,6 +360,7 @@ struct server_slot { spec_draft.clear(); spec_i_batch.clear(); spec_ckpt.clear(); + spec_smpl_save.reset(); } generated_tokens.clear(); generated_token_probs.clear(); @@ -3108,6 +3110,11 @@ private: // update the batch with the sampled/drafted tokens iterate(generating, [&](server_slot & slot) { + GGML_ASSERT(!slot.spec_smpl_save); + if (!slot.spec_draft.empty()) { + // backend sampling advances the sampler during llama_decode() + slot.spec_smpl_save.reset(common_sampler_clone(slot.smpl.get())); + } slot.handle_last_sampled_token(batch); }); @@ -3886,8 +3893,8 @@ private: // verify and try to accept the draft { - // save the sampler sampler state in case we need to restore it - common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); + GGML_ASSERT(slot.spec_smpl_save); + common_sampler_ptr smpl_save = std::move(slot.spec_smpl_save); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 5fc338ad78..12e391eedc 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -25,34 +25,37 @@ def fixture_create_server(): def test_with_and_without_draft(): global server - server.model_draft = None # disable draft model - server.spec_type = None - server.start() - res = server.make_request("POST", "/completion", data={ + request = { "prompt": "I believe the meaning of life is", "temperature": 0.0, "top_k": 1, + "seed": 4242, "n_predict": 16, - }) + "return_tokens": True, + } + + server.model_draft = None # disable draft model + server.spec_type = None + server.backend_sampling = True + server.start() + res = server.make_request("POST", "/completion", data=request) assert res.status_code == 200 content_no_draft = res.body["content"] + tokens_no_draft = res.body["tokens"] server.stop() # 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", - "temperature": 0.0, - "top_k": 1, - "n_predict": 16, - }) + res = server.make_request("POST", "/completion", data=request) assert res.status_code == 200 assert res.body["timings"]["draft_n"] > 0 content_draft = res.body["content"] + tokens_draft = res.body["tokens"] assert content_no_draft == content_draft + assert tokens_no_draft == tokens_draft def test_different_draft_min_draft_max():