Don't reuse memory for output views.

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) {

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<llama_seq_id> 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<uint32_t> sampler_seqs;
+    std::vector<bool> 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;
             }
 

src/llama-sampler.cpp

@@ -1089,8 +1089,6 @@ struct llama_sampler_dist : public llama_sampler_backend {
 
     // inputs for the current sampling graph
     std::vector<ggml_tensor *> 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;

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); });

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;
             }