diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index 3bda9abbe0..65eeaedc42 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -805,11 +805,14 @@ static void ggml_gallocr_alloc_graph_impl(ggml_gallocr_t galloc, struct ggml_cgr if (ggml_impl_is_view(parent)) { struct ggml_tensor * view_src = parent->view_src; struct hash_node * view_src_hn = ggml_gallocr_hash_get(galloc, view_src); - view_src_hn->n_views -= 1; - AT_PRINTF("view_src %s: %d children, %d views\n", - view_src->name, view_src_hn->n_children, view_src_hn->n_views); - if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) { - ggml_gallocr_free_node(galloc, view_src); + // output views keep their source alive until graph completion + if (!(parent->flags & GGML_TENSOR_FLAG_OUTPUT)) { + view_src_hn->n_views -= 1; + AT_PRINTF("view_src %s: %d children, %d views\n", + view_src->name, view_src_hn->n_children, view_src_hn->n_views); + if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) { + ggml_gallocr_free_node(galloc, view_src); + } } } else if (p_hn->allocated) { diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ed4bf6a375..170f1a5e58 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2353,7 +2353,7 @@ 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); - // spread the outputs across all sequences to reserve the largest sampling graph + // select sampler outputs first to reserve the largest valid sampling graph std::vector seq_ids(n_seqs); for (uint32_t s = 0; s < n_seqs; ++s) { seq_ids[s] = s; @@ -2365,11 +2365,50 @@ ggml_cgraph * llama_context::graph_reserve( } uint32_t n_outputs_set = 0; + + std::vector sampler_seqs; + std::vector has_sampler(n_seqs, false); + for (const auto & entry : sampling.samplers) { + const llama_seq_id seq_id = entry.first; + if (seq_id < 0 || (uint32_t) seq_id >= n_seqs) { + continue; + } + + sampler_seqs.push_back(seq_id); + has_sampler[seq_id] = true; + } + + 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; + } + + ubatch.output[s * ubatch.n_seq_tokens + t] = true; + ++n_outputs_set; + } + } + + // use sequences without samplers for any remaining outputs 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() || - t >= cparams.n_sampling_outputs_per_seq_max)) { + if (has_sampler[s]) { continue; } diff --git a/src/llama-sampler.cpp b/src/llama-sampler.cpp index 3f82bd768d..e54cdcff07 100644 --- a/src/llama-sampler.cpp +++ b/src/llama-sampler.cpp @@ -1089,8 +1089,6 @@ 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) { @@ -1199,8 +1197,7 @@ 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; - - sctx->copy_candidates = n_outputs_per_seq_max > 1; + GGML_UNUSED(n_outputs_per_seq_max); const bool res = llama_sampler_backend_support(smpl, buft); @@ -1267,12 +1264,6 @@ static void llama_sampler_dist_backend_apply( sampled_token = ggml_get_rows(ctx, candidates, idx); ggml_set_name(sampled_token, "dist_sampled_token"); - - 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-alloc.cpp b/tests/test-alloc.cpp index 6d5428493e..ad8dd21c97 100644 --- a/tests/test-alloc.cpp +++ b/tests/test-alloc.cpp @@ -388,6 +388,25 @@ static void test_view_inplace() { GGML_ASSERT(backend.context->allocated_total() <= 24); } +static void test_output_view_lifetime() { + dummy_backend backend = dummy_backend_init(SIZE_MAX); + auto [ctx, graph, ctx_ptr] = make_context(); + + ggml_tensor * x[5]; + x[0] = make_input_1d(ctx, 4); + x[1] = ggml_scale(ctx, x[0], 2.0f); + x[2] = ggml_reshape_2d(ctx, x[1], 2, 2); + x[3] = ggml_sum(ctx, x[2]); + x[4] = ggml_pad(ctx, x[3], 3, 0, 0, 0); + assign_names(ctx); + + ggml_set_output(x[2]); + ggml_gallocr_ptr galloc = allocate_graph(graph, x[4], &backend.buffer_type); + check_all_allocated(graph); + check_max_size(ctx); + GGML_ASSERT(!memory_overlap(x[2], x[4])); +} + static void test_reuse_and_free() { dummy_backend backend = dummy_backend_init(40); auto [ctx, graph, ctx_ptr] = make_context(); @@ -597,6 +616,7 @@ int main() { run("test_not_enough_chunks", test_not_enough_chunks); run("test_fill_leftover_space", test_fill_leftover_space); run("test_view_inplace", test_view_inplace); + run("test_output_view_lifetime", test_output_view_lifetime); run("test_reuse_and_free", test_reuse_and_free); run("test_merge_free_block(32)", []() { test_merge_free_block(32); }); run("test_merge_free_block(SIZE_MAX)", []() { test_merge_free_block(SIZE_MAX); }); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9c6e638b1..aa7f4037a9 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3882,7 +3882,8 @@ private: // speculative decoding - main model sample and accept iterate(slots, [&](server_slot & slot) { - if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || slot.spec_draft.empty()) { + if (slot.state != SLOT_STATE_GENERATING || !slot.can_speculate() || + slot.spec_draft.empty() || slot.spec_i_batch.empty()) { return; }