llama : support multi-output backend sampling (#25532)

* Enable backend sampling with token speculation

* Clamp the mask sum before converting it into the sampled index

* Add a numeric context parameter declaring the maximum outputs one sequence

* More fixes

* Don't reuse memory for output views.

* Match dist between CPU and GPU

* Fix CPU and backend sampling mismatches

* Simpify some of the changes

* Fix tests on Vulkan

* More test fixes

* Rebase changes

* Rebase and address review comments

* Address review comments

* Address review comments

* Update src/llama-sampler.cpp

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

---------

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

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_outputs_max_per_seq = std::max(params.n_outputs_max_per_seq, 0);
     cparams.n_batch           = params.n_batch;
     cparams.n_ubatch          = params.n_ubatch;
     cparams.n_threads         = params.cpuparams.n_threads;

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_outputs_max_per_seq =     1; // max outputs per sequence
     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)

common/llguidance.cpp

@@ -116,6 +116,8 @@ static llama_sampler_i llama_sampler_llg_i = {
     /* .backend_accept    = */ NULL,
     /* .backend_apply     = */ NULL,
     /* .backend_set_input = */ NULL,
+    /* .backend_reset     = */ NULL,
+    /* .copy_state        = */ NULL,
 };
 
 static size_t llama_sampler_llg_tokenize_fn(const void * user_data, const uint8_t * bytes, size_t bytes_len,

common/reasoning-budget.cpp

@@ -217,6 +217,8 @@ static struct llama_sampler_i common_reasoning_budget_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl) {

common/sampling.cpp

@@ -518,6 +518,26 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) {
     };
 }
 
+void common_sampler_copy(const common_sampler * src, common_sampler * dst) {
+    if (!src || !dst || src == dst) {
+        return;
+    }
+
+    GGML_ASSERT((src->grmr == nullptr) == (dst->grmr == nullptr));
+    GGML_ASSERT((src->rbudget == nullptr) == (dst->rbudget == nullptr));
+
+    llama_sampler_copy(src->grmr,    dst->grmr);
+    llama_sampler_copy(src->rbudget, dst->rbudget);
+    llama_sampler_copy(src->chain,   dst->chain);
+
+    dst->params     = src->params;
+    dst->prev       = src->prev;
+    dst->cur        = src->cur;
+    dst->cur_p      = src->cur_p;
+    dst->cur_p.data = src->cur_p.data ? dst->cur.data() : nullptr; // re-point to dst's buffer
+    dst->t_total_us = src->t_total_us;
+}
+
 void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl) {
     // TODO: measure grammar performance
 

common/sampling.h

@@ -47,6 +47,7 @@ void common_sampler_free(struct common_sampler * gsmpl);
 void                    common_sampler_accept(struct common_sampler * gsmpl, llama_token token, bool is_generated);
 void                    common_sampler_reset (struct common_sampler * gsmpl);
 struct common_sampler * common_sampler_clone (struct common_sampler * gsmpl);
+void                    common_sampler_copy  (const struct common_sampler * src, struct common_sampler * dst);
 
 // arguments can be nullptr to skip printing
 void common_perf_print(const struct llama_context * ctx, const struct common_sampler * gsmpl);

common/speculative.cpp

@@ -2299,6 +2299,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_outputs_max_per_seq = 1;
 
     return result;
 }
@@ -2384,6 +2385,17 @@ common_speculative_init_result_ptr common_speculative_init_from_params(common_pa
     return std::make_unique<common_speculative_init_result>(params, model_tgt, ctx_tgt);
 }
 
+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;
+
+    return {
+        /* .total   = */ (int32_t) std::min<int64_t>(n_batch, total),
+        /* .per_seq = */ (int32_t) std::min<int64_t>(n_batch, per_seq),
+    };
+}
+
 // initialization of the speculative decoding system
 //
 common_speculative * common_speculative_init(common_params_speculative & params, uint32_t n_seq) {

common/speculative.h

@@ -25,6 +25,15 @@ int32_t common_speculative_n_max(const common_params_speculative * spec);
 
 common_params common_base_params_to_speculative(const common_params & params);
 
+struct common_speculative_output_limits {
+    int32_t total;
+    int32_t per_seq;
+};
+
+// 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);
 
 void common_speculative_free(common_speculative * spec);

docs/speculative.md

@@ -202,6 +202,12 @@ Example Video:
 
 If a draft model is combined with a draftless decoding the draftless decoding has higher precedence.
 
+### Backend Sampling
+
+Use `--backend-sampling` to run supported target-model samplers on the model backend. Draft-model sampling uses the backend by default and can be controlled with `--spec-draft-backend-sampling` and `--no-spec-draft-backend-sampling`.
+
+Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required.
+
 ### General Speculative Parameters
 
 ```

examples/lookup/lookup.cpp

@@ -3,9 +3,11 @@
 #include "common.h"
 #include "ngram-cache.h"
 #include "sampling.h"
+#include "speculative.h"
 #include "log.h"
 #include "llama.h"
 
+#include <algorithm>
 #include <clocale>
 #include <cstdint>
 #include <cstdio>
@@ -27,6 +29,10 @@ 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;
 
+    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_outputs_max_per_seq = output_limits.per_seq;
+
     // init llama.cpp
     llama_backend_init();
     llama_numa_init(params.numa);

examples/speculative-simple/speculative-simple.cpp

@@ -5,6 +5,7 @@
 #include "log.h"
 #include "llama.h"
 
+#include <algorithm>
 #include <clocale>
 #include <cstdio>
 #include <cstring>
@@ -29,6 +30,11 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    const auto output_limits = common_speculative_get_output_limits(
+            params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
+    params.n_outputs_max = output_limits.total;
+    params.n_outputs_max_per_seq = output_limits.per_seq;
+
     // init llama.cpp
     llama_backend_init();
     llama_numa_init(params.numa);
@@ -55,6 +61,9 @@ int main(int argc, char ** argv) {
 
         auto params_dft = params;
 
+        params_dft.n_outputs_max = params.n_parallel;
+        params_dft.n_outputs_max_per_seq = 1;
+
         params_dft.devices      = params_spec.devices;
         params_dft.model        = params_spec.mparams;
         params_dft.n_gpu_layers = params_spec.n_gpu_layers;

examples/speculative/speculative.cpp

@@ -1,6 +1,7 @@
 #include "arg.h"
 #include "common.h"
 #include "sampling.h"
+#include "speculative.h"
 #include "log.h"
 #include "llama.h"
 
@@ -57,6 +58,11 @@ int main(int argc, char ** argv) {
     // max number of parallel drafting sequences (i.e. tree branches)
     const int n_seq_dft = params.n_parallel;
 
+    const auto output_limits = common_speculative_get_output_limits(
+            params.n_batch, params.n_parallel, params.speculative.draft.n_max);
+    params.n_outputs_max = output_limits.total;
+    params.n_outputs_max_per_seq = 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;
 
@@ -83,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_outputs_max_per_seq = 1;
     if (params.speculative.draft.cpuparams.n_threads > 0) {
         params.cpuparams.n_threads = params.speculative.draft.cpuparams.n_threads;
     }

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_outputs_max_per_seq; // max outputs per sequence (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
 
@@ -1054,6 +1055,9 @@ 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);
 
@@ -1270,9 +1274,12 @@ extern "C" {
         // [EXPERIMENTAL]
         // backend sampling interface:
 
-        // return true if the backend supports all ops needed by the sampler
+        // return true if the backend supports all ops needed by the sampler and can handle up to n_outputs_max_per_seq outputs per sequence
         // note: call once per sampler
-        bool (*backend_init)(struct llama_sampler * smpl, ggml_backend_buffer_type_t buft);
+        bool (*backend_init)(
+                struct llama_sampler       * smpl,
+                ggml_backend_buffer_type_t   buft,
+                uint32_t                     n_outputs_max_per_seq);
 
         // call after .backend_apply()
         void (*backend_accept)(
@@ -1290,6 +1297,13 @@ extern "C" {
 
         // 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 any internal sampler state
+        void (*backend_reset)(struct llama_sampler * smpl);
+
+        // copy mutable state from src into dst while keeping dst's references to the current sampling graph
+        // src and dst must have the same type and configuration
+        void (*copy_state)(const struct llama_sampler * src, struct llama_sampler * dst);
     };
 
     struct llama_sampler {
@@ -1310,6 +1324,9 @@ extern "C" {
     LLAMA_API void                   llama_sampler_apply (      struct llama_sampler * smpl, llama_token_data_array * cur_p);
     LLAMA_API void                   llama_sampler_reset (      struct llama_sampler * smpl);
     LLAMA_API struct llama_sampler * llama_sampler_clone (const struct llama_sampler * smpl);
+    // copy mutable sampler state without changing dst or its sampling graph bindings
+    // src and dst must have the same type and configuration
+    LLAMA_API void                   llama_sampler_copy  (const struct llama_sampler * src, struct llama_sampler * dst);
     // important: do not free if the sampler has been added to a llama_sampler_chain (via llama_sampler_chain_add)
     LLAMA_API void                   llama_sampler_free  (      struct llama_sampler * smpl);
 
@@ -1499,6 +1516,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);

src/llama-context.cpp

@@ -10,6 +10,7 @@
 #include "llama-mmap.h"
 #include "llama-model.h"
 #include "llama-ext.h"
+#include "llama-sampler.h"
 #include "llama.h"
 
 #include <cinttypes>
@@ -159,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;
@@ -265,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_outputs_max_per_seq = params.n_outputs_max_per_seq == 0 ?
+            cparams.n_outputs_max : std::min(params.n_outputs_max_per_seq, 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;
@@ -312,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_outputs_max_per_seq = %u\n",   __func__, cparams.n_outputs_max_per_seq);
 
     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",
@@ -1231,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);
+        sampler->iface->backend_init(sampler, buft, cparams.n_outputs_max_per_seq);
 
         sampling.samplers[seq_id] = sampler;
 
@@ -1576,109 +1580,39 @@ int llama_context::encode(const llama_batch & batch_inp) {
     return 0;
 }
 
-static std::map<llama_seq_id, uint32_t> build_seq_to_output_row(const llama_ubatch & ubatch, uint32_t row_offset) {
+template<typename T>
-    std::map<llama_seq_id, uint32_t> seq_to_row;
+static void copy_tensor_async_rows(
-    // how many output tokens we have seen so far for this ubatch.
+    const std::vector<ggml_tensor *> & tensors,
-    uint32_t local = 0;
+    const buffer_view<T> & dst,
-    for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
-        // skip tokens that are not output.
-        if (!ubatch.output[i]) {
-            continue;
-        }
-
-        const llama_seq_id seq_id = ubatch.seq_id[i][0];
-        // row_offset is the number of output tokens before this ubatch.
-        seq_to_row[seq_id] = row_offset + local;
-        ++local;
-    }
-    return seq_to_row;
-}
-
-static void copy_tensor_async_ints(
-    const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
-    const buffer_view<llama_token> & sampled,
-    const std::map<llama_seq_id, uint32_t> & seq_to_row,
-    ggml_backend_sched_t sched) {
-    if (!sampled.has_data()) {
-        return;
-    }
-
-    for (const auto & [seq_id, tensor] : tensor_map) {
-        auto it = seq_to_row.find(seq_id);
-        if (it == seq_to_row.end()) {
-            continue;
-        }
-
-        const uint32_t row = it->second;
-        GGML_ASSERT(row < sampled.size);
-
-        GGML_ASSERT(ggml_is_contiguous(tensor) && "sampled tokens tensor must be contiguous for async copy");
-
-        ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
-        ggml_backend_tensor_get_async(backend, tensor, sampled.data + row, 0, sizeof(sampled.data[row]));
-    }
-}
-
-static void copy_tensor_async_floats(
-    const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
-    const buffer_view<float> & dst,
     size_t stride,
-    std::vector<uint32_t> & counts,
+    uint32_t row_offset,
-    const std::map<llama_seq_id, uint32_t> & seq_to_row,
+    ggml_backend_sched_t sched,
-    ggml_backend_sched_t sched) {
+    std::vector<uint32_t> * counts = nullptr) {
     if (!dst.has_data()) {
         return;
     }
 
-    for (const auto & [seq_id, tensor] : tensor_map) {
+    for (size_t i = 0; i < tensors.size(); ++i) {
-        auto it = seq_to_row.find(seq_id);
+        auto * tensor = tensors[i];
-        if (it == seq_to_row.end()) {
+        if (tensor == nullptr) {
             continue;
         }
 
-        const uint32_t row = it->second;
+        const uint32_t row = row_offset + i;
-        GGML_ASSERT(row < counts.size());
+        const size_t n_elements = ggml_nelements(tensor);
-
+        GGML_ASSERT(ggml_is_contiguous(tensor) && "sampling tensor must be contiguous for async copy");
-        GGML_ASSERT(ggml_is_contiguous(tensor) && "logits/probs tensor must be contiguous for async copy");
+        GGML_ASSERT(n_elements <= stride);
+        GGML_ASSERT((size_t) row * stride + n_elements <= dst.size);
 
         ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
-        float * row_ptr = dst.data + (size_t) row * stride;
+        T * row_ptr = dst.data + (size_t) row * stride;
         ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
 
-        // Update the actual number of logits/probabilities that were written for this row.
+        if (counts) {
-        counts[row] = ggml_nelements(tensor);
+            GGML_ASSERT(row < counts->size());
+            (*counts)[row] = n_elements;
         }
     }
-
-static void copy_tensor_async_candidates(
-    const std::map<llama_seq_id, ggml_tensor*> & tensor_map,
-    const buffer_view<llama_token> & dst,
-    size_t stride,
-    std::vector<uint32_t> & counts,
-    const std::map<llama_seq_id, uint32_t> & seq_to_row,
-    ggml_backend_sched_t sched) {
-    if (!dst.has_data()) {
-        return;
-    }
-
-    for (const auto & [seq_id, tensor] : tensor_map) {
-        auto it = seq_to_row.find(seq_id);
-        if (it == seq_to_row.end()) {
-            continue;
-        }
-
-        const uint32_t row = it->second;
-        GGML_ASSERT(row < counts.size());
-
-        GGML_ASSERT(ggml_is_contiguous(tensor) && "candidates tensor must be contiguous for async copy");
-
-        ggml_backend_t backend = ggml_backend_sched_get_tensor_backend(sched, tensor);
-        llama_token * row_ptr = dst.data + (size_t) row * stride;
-        ggml_backend_tensor_get_async(backend, tensor, row_ptr, 0, ggml_nbytes(tensor));
-
-        // Update the actual number of candidates that were written.
-        counts[row] = ggml_nelements(tensor);
-    }
 }
 
 static bool needs_raw_logits(const llama_ubatch & ubatch, const std::map<llama_seq_id, llama_sampler *> & samplers) {
@@ -1726,12 +1660,12 @@ int llama_context::decode(const llama_batch & batch_inp) {
 
     const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max;
 
-    // TODO: avoid this workaround in the future
+    // embedding contexts output every token even when batch.logits is not set
-    if (has_samplers && batch_inp.logits) {
+    if (has_samplers && (output_all || batch_inp.logits)) {
         std::vector<int32_t> seq_output_count(n_seq_max, 0);
 
         for (int32_t i = 0; i < batch_inp.n_tokens; ++i) {
-            if (batch_inp.logits[i] == 0) {
+            if (!output_all && batch_inp.logits[i] == 0) {
                 continue;
             }
 
@@ -1740,10 +1674,17 @@ 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]++;
-                if (seq_output_count[seq_id] > 1) {
+                auto sampler = sampling.samplers.find(seq_id);
-                    LLAMA_LOG_ERROR("%s: backend sampling requires at most one output token per sequence (seq_id %d had %d)\n",
+                if (sampler != sampling.samplers.end() &&
-                            __func__, seq_id, seq_output_count[seq_id]);
+                        seq_output_count[seq_id] > (int32_t) cparams.n_outputs_max_per_seq) {
+                    LLAMA_LOG_ERROR("%s: backend sampling supports at most %u outputs per sequence "
+                            "(seq_id %d had %d)\n", __func__, cparams.n_outputs_max_per_seq,
+                            seq_id, seq_output_count[seq_id]);
                     return -1;
                 }
             }
@@ -1843,6 +1784,11 @@ int llama_context::decode(const llama_batch & batch_inp) {
         return -2;
     };
 
+    // start a new sampling transaction for this logical batch
+    for (const auto & entry : sampling.samplers) {
+        llama_sampler_backend_begin(entry.second);
+    }
+
     int64_t n_outputs_prev = 0;
     int64_t n_tokens_prev  = 0;
 
@@ -2009,17 +1955,14 @@ int llama_context::decode(const llama_batch & batch_inp) {
             }
         }
 
-        // Copy backend sampling output if this ubatch produced any sampling tensors.
+        if (has_samplers) {
-        if (has_samplers && (!res->t_sampled.empty() || !res->t_sampled_probs.empty() || !res->t_sampled_logits.empty())) {
-            const auto seq_to_output_row = build_seq_to_output_row(ubatch, n_outputs_prev);
             const auto stride = n_vocab;
 
             // async copy the sampling data from the backend to the host
-            copy_tensor_async_ints(res->t_sampled, sampling.sampled, seq_to_output_row, sched.get());
+            copy_tensor_async_rows(res->t_sampled,        sampling.sampled,    1,      n_outputs_prev, sched.get());
-
+            copy_tensor_async_rows(res->t_sampled_logits, sampling.logits,     stride, n_outputs_prev, sched.get(), &sampling.logits_count);
-            copy_tensor_async_floats    (res->t_sampled_logits, sampling.logits,     stride, sampling.logits_count,     seq_to_output_row, sched.get());
+            copy_tensor_async_rows(res->t_sampled_probs,  sampling.probs,      stride, n_outputs_prev, sched.get(), &sampling.probs_count);
-            copy_tensor_async_floats    (res->t_sampled_probs,  sampling.probs,      stride, sampling.probs_count,      seq_to_output_row, sched.get());
+            copy_tensor_async_rows(res->t_candidates,     sampling.candidates, stride, n_outputs_prev, sched.get(), &sampling.candidates_count);
-            copy_tensor_async_candidates(res->t_candidates,     sampling.candidates, stride, sampling.candidates_count, seq_to_output_row, sched.get());
         }
 
         n_outputs_prev += n_outputs;
@@ -2349,6 +2292,7 @@ void llama_context::output_reorder() {
 //
 
 uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
+    uint32_t res;
     if (model.arch == LLM_ARCH_QWEN3NEXT ||
         model.arch == LLM_ARCH_KIMI_LINEAR ||
         model.arch == LLM_ARCH_QWEN35 ||
@@ -2357,12 +2301,32 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
         (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) ||
         model.arch == LLM_ARCH_NANBEIGE ||
         model.arch == LLM_ARCH_MINIMAX_M3) {
-        return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
+        res = std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors());
-    }
+    } else {
-    uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
+        res = std::max<uint32_t>(1024u, 8u*model.n_tensors());
         for (const auto & lora : model.loras) {
             res += lora->get_n_nodes();
         }
+    }
+
+    uint32_t n_sampling_nodes = 0;
+    uint32_t n_sampling_nodes_max = 0;
+    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.n_outputs_max_per_seq > 1) {
+            n_sampling_nodes_max = std::max(n_sampling_nodes_max, n_nodes);
+        }
+    }
+
+    const uint32_t n_sampling_outputs_max = std::min<uint64_t>(
+            std::min(n_tokens, cparams.n_outputs_max),
+            (uint64_t) cparams.n_seq_max * cparams.n_outputs_max_per_seq);
+
+    res += n_sampling_nodes;
+    if (n_sampling_outputs_max > 1) {
+        res += (n_sampling_outputs_max - 1) * n_sampling_nodes_max;
+    }
     return res;
 }
 
@@ -2370,6 +2334,63 @@ llm_graph_result * llama_context::get_gf_res_reserve() const {
     return static_cast<llm_graph_result *>(gf_res_reserve.get());
 }
 
+// pack sampler outputs into as few sequences as possible before using sequences without samplers
+static void ubatch_prepare_reserve(
+              llama_ubatch                            & ubatch,
+              uint32_t                                  n_outputs,
+        const std::map<llama_seq_id, llama_sampler *> & samplers,
+              uint32_t                                  n_outputs_max_per_seq) {
+    const uint32_t n_seqs       = ubatch.n_seqs;
+    const uint32_t n_seq_tokens = ubatch.n_seq_tokens;
+
+    for (uint32_t s = 0; s < n_seqs; ++s) {
+        for (uint32_t t = 0; t < n_seq_tokens; ++t) {
+            const uint32_t i = s * n_seq_tokens + t;
+            ubatch.n_seq_id[i] = 1;
+            ubatch.seq_id[i] = &ubatch.seq_id_unq[s];
+        }
+    }
+
+    // sequences with a sampler that fit in this ubatch
+    std::vector<uint32_t> sampler_seqs;
+    std::vector<bool> has_sampler(n_seqs, false);
+    for (const auto & entry : 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;
+    }
+
+    uint32_t n_outputs_set = 0;
+
+    const uint32_t n_outputs_per_seq = std::min(n_seq_tokens, n_outputs_max_per_seq);
+    for (uint32_t s : sampler_seqs) {
+        if (n_outputs_set >= n_outputs) {
+            break;
+        }
+
+        for (uint32_t t = 0; t < n_outputs_per_seq && n_outputs_set < n_outputs; ++t) {
+            ubatch.output[s * n_seq_tokens + t] = true;
+            ++n_outputs_set;
+        }
+    }
+
+    // use sequences without samplers for any remaining outputs
+    for (uint32_t t = 0; t < n_seq_tokens && n_outputs_set < n_outputs; ++t) {
+        for (uint32_t s = 0; s < n_seqs && n_outputs_set < n_outputs; ++s) {
+            if (has_sampler[s]) {
+                continue;
+            }
+
+            ubatch.output[s * n_seq_tokens + t] = true;
+            ++n_outputs_set;
+        }
+    }
+}
+
 ggml_cgraph * llama_context::graph_reserve(
         uint32_t n_tokens, uint32_t n_seqs, uint32_t n_outputs, const llama_memory_context_i * mctx, bool split_only, size_t * sizes) {
     LLAMA_LOG_DEBUG("%s: reserving a graph for ubatch with n_tokens = %4u, n_seqs = %2u, n_outputs = %4u\n", __func__, n_tokens, n_seqs, n_outputs);
@@ -2394,14 +2415,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);
 
-    // set one output token per sequence in order to activate all backend samplers
+    ubatch_prepare_reserve(ubatch, n_outputs, sampling.samplers, cparams.n_outputs_max_per_seq);
-    std::vector<llama_seq_id> seq_ids(n_seqs);
-    for (uint32_t i = 0; i < n_seqs; ++i) {
-        seq_ids[i] = i;
-        ubatch.n_seq_id[i] = 1;
-        ubatch.seq_id[i] = &seq_ids[i];
-        ubatch.output[i] = true;
-    }
 
     auto * res = gf_res_reserve.get();
 
@@ -3488,6 +3502,7 @@ llama_context_params llama_context_default_params() {
         /*.n_seq_max                   =*/ 1,
         /*.n_rs_seq                    =*/ 0,
         /*.n_outputs_max               =*/ 0,
+        /*.n_outputs_max_per_seq       =*/ 1,
         /*.n_threads                   =*/ GGML_DEFAULT_N_THREADS, // TODO: better default
         /*.n_threads_batch             =*/ GGML_DEFAULT_N_THREADS,
         /*.ctx_type                    =*/ LLAMA_CONTEXT_TYPE_DEFAULT,

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_outputs_max_per_seq;
     int32_t  n_threads;       // number of threads to use for generation
     int32_t  n_threads_batch; // number of threads to use for batch processing
 

src/llama-graph.cpp

@@ -4,6 +4,7 @@
 #include "llama-model.h"
 #include "llama-batch.h"
 #include "llama-cparams.h"
+#include "llama-sampler.h"
 
 #include "llama-kv-cache.h"
 #include "llama-kv-cache-iswa.h"
@@ -1353,24 +1354,24 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) {
             }
         }
     }
-    for (auto & [seq_id, t] : t_sampled) {
+    for (auto * tensor : t_sampled) {
-        if (t != nullptr) {
+        if (tensor != nullptr) {
-            ggml_set_output(t);
+            ggml_set_output(tensor);
         }
     }
-    for (auto & [seq_id, t] : t_sampled_probs) {
+    for (auto * tensor : t_sampled_probs) {
-        if (t != nullptr) {
+        if (tensor != nullptr) {
-            ggml_set_output(t);
+            ggml_set_output(tensor);
         }
     }
-    for (auto & [seq_id, t] : t_sampled_logits) {
+    for (auto * tensor : t_sampled_logits) {
-        if (t != nullptr) {
+        if (tensor != nullptr) {
-            ggml_set_output(t);
+            ggml_set_output(tensor);
         }
     }
-    for (auto & [seq_id, t] : t_candidates) {
+    for (auto * tensor : t_candidates) {
-        if (t != nullptr) {
+        if (tensor != nullptr) {
-            ggml_set_output(t);
+            ggml_set_output(tensor);
         }
     }
 }
@@ -3649,34 +3650,45 @@ void llm_graph_context::build_sampling() const {
     auto inp_sampling = std::make_unique<llm_graph_input_sampling>(samplers);
     res->add_input(std::move(inp_sampling));
 
-    std::map<llama_seq_id, int32_t> seq_to_logit_row;
+    std::map<llama_seq_id, std::vector<uint32_t>> sampling_rows;
-    int32_t logit_row_idx = 0;
+    uint32_t n_rows = 0;
-
+    for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
-    for (uint32_t i = 0; i < ubatch.n_tokens; i++) {
         if (ubatch.output[i]) {
-            llama_seq_id seq_id = ubatch.seq_id[i][0];
+            sampling_rows[ubatch.seq_id[i][0]].push_back(n_rows++);
-            seq_to_logit_row[seq_id] = logit_row_idx;
-            logit_row_idx++;
         }
     }
 
+    res->t_sampled.resize(n_rows, nullptr);
+    res->t_sampled_probs.resize(n_rows, nullptr);
+    res->t_sampled_logits.resize(n_rows, nullptr);
+    res->t_candidates.resize(n_rows, nullptr);
+
     // res->t_logits will contain logits for all tokens that want the logits calculated (logits=1 or output=1)
     GGML_ASSERT(res->t_logits != nullptr && "missing t_logits tensor");
 
-    // add a dummy row of logits
+    // add a dummy row to keep the single-output graph static regardless of active samplers
-    // this trick makes the graph static, regardless of which samplers are activated
+    // multi-output graphs can still vary with the number of output rows
-    // this is important in order to minimize graph reallocations
     ggml_tensor * logits_t = ggml_pad(ctx0, res->t_logits, 0, 1, 0, 0);
 
+    for (const auto & entry : samplers) {
+        if (entry.second->iface->backend_reset) {
+            entry.second->iface->backend_reset(entry.second);
+        }
+    }
+
+    static const std::vector<uint32_t> dummy_row = { 0 };
+
     for (const auto & [seq_id, sampler] : samplers) {
-        const auto it = seq_to_logit_row.find(seq_id);
+        const auto it = sampling_rows.find(seq_id);
 
         // inactive samplers always work on the first row
-        const auto row_idx = it != seq_to_logit_row.end() ? it->second : 0;
+        const bool active = it != sampling_rows.end();
-        const int i_out    = it != seq_to_logit_row.end() ? 1          : 0;
+        const auto & rows = active ? it->second : dummy_row;
+        const int i_out   = active ? 1          : 0;
 
-        ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], row_idx * logits_t->nb[1]);
+        for (uint32_t i = 0; i < rows.size(); ++i) {
-        ggml_format_name(logits_seq, "logits_seq_%d", seq_id);
+            ggml_tensor * logits_seq = ggml_view_1d(ctx0, logits_t, logits_t->ne[0], rows[i] * logits_t->nb[1]);
+            ggml_format_name(logits_seq, "logits_seq_%d_%u", seq_id, i);
 
             struct llama_sampler_data data = {
                 /*.logits       =*/ logits_seq,
@@ -3689,37 +3701,51 @@ void llm_graph_context::build_sampling() const {
             sampler->iface->backend_apply(sampler, ctx0, gf, &data);
 
             if (data.sampled != nullptr) {
-            res->t_sampled[seq_id] = data.sampled;
+                if (active) {
+                    res->t_sampled[rows[i]] = data.sampled;
+                }
                 outs[1] = data.sampled;
                 ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
             }
 
             if (data.probs != nullptr) {
-            res->t_sampled_probs[seq_id] = data.probs;
+                if (active) {
+                    res->t_sampled_probs[rows[i]] = data.probs;
+                }
                 outs[1] = data.probs;
                 ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
             }
 
             if (data.logits != nullptr) {
-            res->t_sampled_logits[seq_id] = data.logits;
+                if (active) {
+                    res->t_sampled_logits[rows[i]] = data.logits;
+                }
                 outs[1] = data.logits;
                 ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
             }
 
             if (data.candidates != nullptr) {
-            res->t_candidates[seq_id] = data.candidates;
+                if (active) {
+                    res->t_candidates[rows[i]] = data.candidates;
+                }
                 outs[1] = data.candidates;
                 ggml_build_forward_select(gf, outs.data(), outs.size(), i_out);
             }
         }
+    }
 
-    // TODO: Call llama_sampler_accept_ggml after all samplers have been applied.
+    // TODO: Call backend_accept after all samplers have been applied.
     /*
     for (const auto & [seq_id, sampler] : samplers) {
-        if (auto it = res->t_sampled.find(seq_id); it != res->t_sampled.end()) {
+        const auto it = sampling_rows.find(seq_id);
-            ggml_tensor * selected_token = it->second;
+        if (it == sampling_rows.end()) {
-            if (selected_token != nullptr) {
+            continue;
-                llama_sampler_accept_ggml(sampler, ctx0, gf, selected_token);
+        }
+
+        for (uint32_t row : it->second) {
+            ggml_tensor * selected_token = res->t_sampled[row];
+            if (selected_token != nullptr && sampler->iface->backend_accept) {
+                sampler->iface->backend_accept(sampler, ctx0, gf, selected_token);
             }
         }
     }

src/llama-graph.h

@@ -904,10 +904,10 @@ public:
 
     std::vector<ggml_tensor *> t_layer_inp;
 
-    std::map<llama_seq_id, ggml_tensor *> t_sampled_logits;
+    std::vector<ggml_tensor *> t_sampled;
-    std::map<llama_seq_id, ggml_tensor *> t_candidates;
+    std::vector<ggml_tensor *> t_sampled_probs;
-    std::map<llama_seq_id, ggml_tensor *> t_sampled;
+    std::vector<ggml_tensor *> t_sampled_logits;
-    std::map<llama_seq_id, ggml_tensor *> t_sampled_probs;
+    std::vector<ggml_tensor *> t_candidates;
 
     std::vector<llm_graph_input_ptr> inputs;
     std::vector<llm_graph_fused_node> fused_nodes;

src/llama-sampler.cpp

@@ -467,9 +467,11 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     GGML_UNUSED(smpl);
     GGML_UNUSED(buft);
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     return true;
 }
@@ -511,6 +513,8 @@ static struct llama_sampler_i llama_sampler_empty_i = {
     /* .backend_accept    = */ llama_sampler_empty_backend_accept,
     /* .backend_apply     = */ llama_sampler_empty_backend_apply,
     /* .backend_set_input = */ llama_sampler_empty_backend_set_input,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_empty(const char * name) {
@@ -551,6 +555,12 @@ struct llama_sampler_backend {
         this->support = support;
     }
 
+    // copy the state that is not tied to the current sampling graph
+    // samplers that hold only immutable configuration can use this as is
+    void copy_state(const llama_sampler_backend & src) {
+        GGML_UNUSED(src);
+    }
+
 private:
     std::string name;
     std::string name_ext;
@@ -559,6 +569,71 @@ private:
     bool support;
 };
 
+// .copy_state for samplers deriving from llama_sampler_backend
+template<typename T>
+static void llama_sampler_backend_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {
+    ((T *) dst->ctx)->copy_state(*(const T *) src->ctx);
+}
+
+struct llama_sampler_backend_probe {
+    ggml_context_ptr ctx;
+    ggml_cgraph * gf;
+};
+
+static llama_sampler_backend_probe llama_sampler_backend_probe_graph(
+        llama_sampler * sampler,
+        int64_t         n_candidates,
+        uint32_t        max_nodes,
+        bool            with_candidates) {
+    ggml_init_params params = {
+        /*.mem_size   =*/ max_nodes * ggml_tensor_overhead() + ggml_graph_overhead_custom(max_nodes, false),
+        /*.mem_buffer =*/ nullptr,
+        /*.no_alloc   =*/ true,
+    };
+
+    ggml_context_ptr ctx_ptr { ggml_init(params) };
+    if (!ctx_ptr) {
+        throw std::runtime_error(format("failed to create ggml context"));
+    }
+
+    auto * ctx = ctx_ptr.get();
+    auto * gf = ggml_new_graph_custom(ctx, max_nodes, false);
+
+    llama_sampler_data data = {
+        /*.logits       =*/ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_candidates),
+        /*.probs        =*/ nullptr,
+        /*.sampled      =*/ nullptr,
+        /*.candidates   =*/ with_candidates ? ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n_candidates) : nullptr,
+    };
+
+    if (sampler->iface->backend_reset) {
+        sampler->iface->backend_reset(sampler);
+    }
+    sampler->iface->backend_apply(sampler, ctx, gf, &data);
+
+    for (auto * output : { data.logits, data.probs, data.sampled, data.candidates }) {
+        if (output) {
+            ggml_build_forward_expand(gf, output);
+        }
+    }
+
+    if (sampler->iface->backend_reset) {
+        sampler->iface->backend_reset(sampler);
+    }
+
+    return { std::move(ctx_ptr), gf };
+}
+
+static uint32_t llama_sampler_backend_probe_n_nodes(const llama_sampler_backend_probe & probe) {
+    uint32_t n_tensors = 0;
+    for (auto * tensor = ggml_get_first_tensor(probe.ctx.get()); tensor;
+            tensor = ggml_get_next_tensor(probe.ctx.get(), tensor)) {
+        ++n_tensors;
+    }
+
+    return std::max<uint32_t>(ggml_graph_n_nodes(probe.gf), n_tensors);
+}
+
 // check if all ggml ops used by the sampler are supported by the backend
 static bool llama_sampler_backend_support(
         llama_sampler              * smpl,
@@ -569,50 +644,10 @@ static bool llama_sampler_backend_support(
         return true;
     }
 
-    ggml_init_params params = {
+    auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, true);
-        /*.mem_size   =*/ 128*ggml_tensor_overhead() + ggml_graph_overhead(),
-        /*.mem_buffer =*/ NULL,
-        /*.no_alloc   =*/ true,
-    };
 
-    ggml_context_ptr ctx_ptr { ggml_init(params) };
+    for (int i = 0; i < ggml_graph_n_nodes(probe.gf); i++) {
-    if (!ctx_ptr) {
+        struct ggml_tensor * op = ggml_graph_node(probe.gf, i);
-        throw std::runtime_error(format("failed to create ggml context"));
-    }
-
-    ggml_context * ctx = ctx_ptr.get();
-
-    const int64_t n = 1024*1024;
-
-    llama_sampler_data data = {
-        /*.logits     = */ ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n),
-        /*.probs      = */ nullptr,
-        /*.sampled    = */ nullptr,
-        /*.candidates = */ ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n),
-    };
-
-    ggml_cgraph * gf = ggml_new_graph(ctx);
-
-    smpl->iface->backend_apply(smpl, ctx, gf, &data);
-
-    if (data.logits) {
-        ggml_build_forward_expand(gf, data.logits);
-    }
-
-    if (data.probs) {
-        ggml_build_forward_expand(gf, data.probs);
-    }
-
-    if (data.sampled) {
-        ggml_build_forward_expand(gf, data.sampled);
-    }
-
-    if (data.candidates) {
-        ggml_build_forward_expand(gf, data.candidates);
-    }
-
-    for (int i = 0; i < ggml_graph_n_nodes(gf); i++) {
-        struct ggml_tensor * op = ggml_graph_node(gf, i);
 
         if (!ggml_backend_dev_supports_op(device, op)) {
             LLAMA_LOG_WARN("%s: device '%s' does not have support for op %s needed for sampler '%s'\n",
@@ -697,7 +732,8 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * chain = (llama_sampler_chain *) smpl->ctx;
 
     GGML_ASSERT(chain->is_init == false && "llama_sampler_chain_backend_init() called twice");
@@ -705,26 +741,32 @@ static bool llama_sampler_chain_backend_init(
     chain->is_init = true;
 
     bool res = true;
+    bool backend_prefix = true;
 
     for (auto & smpl : chain->samplers) {
-        bool res_cur = true;
+        bool cur_prefix = backend_prefix;
 
         // to be able to run a sampler on the backend, it has to:
         // - have the .backend_init() API implemented
         // - return true during .backend_init()
-        if (smpl.ptr->iface->backend_init) {
+        // - support the requested per-sequence output limit
-            if (!smpl.ptr->iface->backend_init(smpl.ptr, buft)) {
+        if (cur_prefix && smpl.ptr->iface->backend_init) {
-                res_cur = false;
+            if (!smpl.ptr->iface->backend_init(smpl.ptr, buft, n_outputs_max_per_seq)) {
+                cur_prefix = false;
             }
         } else {
-            res_cur = false;
+            cur_prefix = false;
         }
 
-        smpl.is_backend = res_cur;
+        smpl.is_backend = cur_prefix;
+        backend_prefix = cur_prefix;
 
-        res = res && res_cur;
+        res = res && cur_prefix;
     }
 
+    auto probe = llama_sampler_backend_probe_graph(smpl, 1024*1024, GGML_DEFAULT_GRAPH_SIZE, false);
+    chain->n_nodes = llama_sampler_backend_probe_n_nodes(probe);
+
     return res;
 }
 
@@ -780,6 +822,36 @@ static void llama_sampler_chain_backend_set_input(struct llama_sampler * smpl) {
     }
 }
 
+static void llama_sampler_chain_backend_reset(struct llama_sampler * smpl) {
+    auto * chain = (llama_sampler_chain *) smpl->ctx;
+
+    for (auto & entry : chain->samplers) {
+        if (!entry.is_backend) {
+            break;
+        }
+        if (entry.ptr->iface->backend_reset) {
+            entry.ptr->iface->backend_reset(entry.ptr);
+        }
+    }
+}
+
+static void llama_sampler_chain_copy_state(const struct llama_sampler * src, struct llama_sampler * dst) {
+    const auto * src_chain = (const llama_sampler_chain *) src->ctx;
+    auto * dst_chain = (llama_sampler_chain *) dst->ctx;
+
+    GGML_ASSERT(src_chain->samplers.size() == dst_chain->samplers.size());
+
+    for (size_t i = 0; i < src_chain->samplers.size(); ++i) {
+        llama_sampler_copy(src_chain->samplers[i].ptr, dst_chain->samplers[i].ptr);
+    }
+
+    // note: is_init, n_nodes and is_backend belong to the current sampling graph
+    dst_chain->params      = src_chain->params;
+    dst_chain->cur         = src_chain->cur;
+    dst_chain->t_sample_us = src_chain->t_sample_us;
+    dst_chain->n_sample    = src_chain->n_sample;
+}
+
 static struct llama_sampler_i llama_sampler_chain_i = {
     /* .name              = */ llama_sampler_chain_name,
     /* .accept            = */ llama_sampler_chain_accept,
@@ -791,6 +863,8 @@ static struct llama_sampler_i llama_sampler_chain_i = {
     /* .backend_accept    = */ llama_sampler_chain_backend_accept,
     /* .backend_apply     = */ llama_sampler_chain_backend_apply,
     /* .backend_set_input = */ llama_sampler_chain_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_chain_backend_reset,
+    /* .copy_state        = */ llama_sampler_chain_copy_state,
 };
 
 struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_params params) {
@@ -799,6 +873,7 @@ struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_param
         /* .ctx   = */ new llama_sampler_chain {
             /* .params               = */ params,
             /* .is_init              = */ false,
+            /* .n_nodes              = */ 0,
             /* .samplers             = */ {},
             /* .cur                  = */ {},
             /* .t_sample_us          = */ 0,
@@ -807,6 +882,16 @@ struct llama_sampler * llama_sampler_chain_init(struct llama_sampler_chain_param
     );
 }
 
+uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler) {
+    GGML_ASSERT(sampler != nullptr);
+    GGML_ASSERT(sampler->iface == &llama_sampler_chain_i);
+
+    const auto * chain = (const llama_sampler_chain *) sampler->ctx;
+    GGML_ASSERT(chain->is_init);
+
+    return chain->n_nodes;
+}
+
 llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_context * ctx, int32_t idx) {
     const llama_token   sampled_token  = llama_get_sampled_token_ith     (ctx, idx);
     const float *       sampled_probs  = llama_get_sampled_probs_ith     (ctx, idx);
@@ -816,6 +901,7 @@ llama_token llama_sampler_sample(struct llama_sampler * smpl, struct llama_conte
     // If a backend sampler has already sampled a token, return it.
     if (sampled_token != LLAMA_TOKEN_NULL) {
         LLAMA_LOG_DEBUG("%s: Backend sampler selected token for idx %d. Skipping CPU samplers\n", __func__, idx);
+        llama_sampler_accept(smpl, sampled_token);
         return sampled_token;
     }
 
@@ -975,8 +1061,10 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_greedy *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
@@ -1012,6 +1100,8 @@ static struct llama_sampler_i llama_sampler_greedy_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_greedy_backend_apply,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_greedy>,
 };
 
 struct llama_sampler * llama_sampler_init_greedy() {
@@ -1031,7 +1121,25 @@ struct llama_sampler_dist : public llama_sampler_backend {
 
     std::mt19937 rng;
 
-    ggml_tensor * inp_uniform;
+    // TODO: refactor + fix naming
+    //       https://github.com/ggml-org/llama.cpp/pull/25532/changes#r3749906719
+    // use a temporary RNG for multi-output sampling so rejected tokens do not advance rng
+    bool backend_transactional;
+    std::mt19937 rng_backend;
+    size_t n_backend_draws_generated;
+    size_t n_backend_draws_committed;
+
+    // inputs for the current sampling graph
+    std::vector<ggml_tensor *> inp_uniforms;
+
+    void copy_state(const llama_sampler_dist & src) {
+        // note: inp_uniforms and backend_transactional belong to the current sampling graph
+        seed_cur                  = src.seed_cur;
+        rng                       = src.rng;
+        rng_backend               = src.rng_backend;
+        n_backend_draws_generated = src.n_backend_draws_generated;
+        n_backend_draws_committed = src.n_backend_draws_committed;
+    }
 };
 
 static const char * llama_sampler_dist_name(const struct llama_sampler * smpl) {
@@ -1050,7 +1158,11 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
 
     cur_p->selected = 0;
 
+    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
+
     if (cur_p->size == 1) {
+        // keep the RNG state aligned with backend sampling, which draws once per output
+        dist(ctx->rng);
         cur_p->data[0].p = 1.0f;
         return;
     }
@@ -1075,7 +1187,6 @@ static void llama_sampler_dist_apply(struct llama_sampler * smpl, llama_token_da
     // sample from the obtained probabilities and normalize the probs in a single pass
     // this is ~3x faster on Mac with full gpt-oss vocab than the version below
     //
-    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
     const double rnd = dist(ctx->rng);
 
           double sum_run = 0.0f;
@@ -1115,6 +1226,9 @@ static void llama_sampler_dist_reset(struct llama_sampler * smpl) {
     auto * ctx = (llama_sampler_dist *) smpl->ctx;
     ctx->seed_cur = get_rng_seed(ctx->seed);
     ctx->rng.seed(ctx->seed_cur);
+    ctx->rng_backend = ctx->rng;
+    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) {
@@ -1125,7 +1239,12 @@ static struct llama_sampler * llama_sampler_dist_clone(const struct llama_sample
     {
         auto * result_ctx = (llama_sampler_dist *) result->ctx;
 
+        result_ctx->seed_cur                  = ctx->seed_cur;
         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;
@@ -1137,12 +1256,17 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_dist *) smpl->ctx;
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
     sctx->init(res);
+    sctx->backend_transactional = n_outputs_max_per_seq > 1;
+    sctx->rng_backend = sctx->rng;
+    sctx->n_backend_draws_generated = 0;
+    sctx->n_backend_draws_committed = 0;
 
     return res;
 }
@@ -1156,9 +1280,10 @@ static void llama_sampler_dist_backend_apply(
 
     auto * sctx = (llama_sampler_dist *) smpl->ctx;
 
-    sctx->inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
+    ggml_tensor * inp_uniform = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1);
-    ggml_set_name (sctx->inp_uniform, "uniform");
+    ggml_format_name(inp_uniform, "uniform_%zu", sctx->inp_uniforms.size());
-    ggml_set_input(sctx->inp_uniform);
+    ggml_set_input(inp_uniform);
+    sctx->inp_uniforms.push_back(inp_uniform);
 
     // flatten
     struct ggml_tensor * logits = ggml_reshape_1d(ctx, data->logits, ggml_nelements(data->logits));
@@ -1174,7 +1299,7 @@ static void llama_sampler_dist_backend_apply(
     // Recall that each entry in cumsum is the cumulative probability up to that
     // index so values stay negative while the cumulative total is below the
     // random value, and become zero/positive once the threshold is crossed.
-    struct ggml_tensor * diff = ggml_sub(ctx, cumsum, sctx->inp_uniform);
+    struct ggml_tensor * diff = ggml_sub(ctx, cumsum, inp_uniform);
     ggml_set_name(diff, "dist_cumsum");
 
     // The ggml_step function produces a tensor where entries are 1 if the
@@ -1189,6 +1314,9 @@ static void llama_sampler_dist_backend_apply(
     struct ggml_tensor * idxf = ggml_sum(ctx, mask);
     ggml_set_name(idxf, "dist_index_f32");
 
+    // Clamp to prevent out-of-bounds access when computing the index.
+    idxf = ggml_clamp(ctx, idxf, 1.0f, mask->ne[0]);
+
     // Use ggml_scale_bias to scale the index value by -1 and then add the size
     // of the mask to that value so we get the correct index ((-1 * idxf) + n).
     struct ggml_tensor * idx = ggml_cast(ctx, ggml_scale_bias(ctx, idxf, -1.0f, mask->ne[0]), GGML_TYPE_I32);
@@ -1210,22 +1338,52 @@ static void llama_sampler_dist_backend_apply(
 static void llama_sampler_dist_backend_set_input(struct llama_sampler * smpl) {
     auto * sctx = (llama_sampler_dist *) smpl->ctx;
 
-    GGML_ASSERT(sctx->inp_uniform != nullptr);
+    GGML_ASSERT(!sctx->inp_uniforms.empty());
 
     // We sample in double precision and cast to float to match rnd numbers of
-    // llama_dampler_dist which uses double precision (sampling from
+    // llama_sampler_dist which uses double precision (sampling from
     // std::uniform_real_distribution<double> and
     // std::uniform_real_distribution<float> with same rng will produce
     // different sequences).
     std::uniform_real_distribution<double> dist(0.0f, 1.0f);
-    const float rnd = dist(sctx->rng);
 
-    ggml_backend_tensor_set(sctx->inp_uniform, &rnd, 0, sizeof(float));
+    auto & rng = sctx->backend_transactional ? sctx->rng_backend : sctx->rng;
+
+    for (auto * inp_uniform : sctx->inp_uniforms) {
+        GGML_ASSERT(inp_uniform != nullptr);
+
+        const float rnd = dist(rng);
+        ggml_backend_tensor_set(inp_uniform, &rnd, 0, sizeof(float));
+
+        if (sctx->backend_transactional) {
+            ++sctx->n_backend_draws_generated;
+        }
+    }
+}
+
+static void llama_sampler_dist_backend_reset(struct llama_sampler * smpl) {
+    auto * sctx = (llama_sampler_dist *) smpl->ctx;
+    sctx->inp_uniforms.clear();
+}
+
+static void llama_sampler_dist_accept(struct llama_sampler * smpl, llama_token token) {
+    GGML_UNUSED(token);
+
+    auto * sctx = (llama_sampler_dist *) smpl->ctx;
+
+    if (!sctx->backend_transactional ||
+            sctx->n_backend_draws_committed >= sctx->n_backend_draws_generated) {
+        return;
+    }
+
+    std::uniform_real_distribution<double> dist(0.0f, 1.0f);
+    dist(sctx->rng);
+    ++sctx->n_backend_draws_committed;
 }
 
 static struct llama_sampler_i llama_sampler_dist_i = {
     /* .name              = */ llama_sampler_dist_name,
-    /* .accept            = */ nullptr,
+    /* .accept            = */ llama_sampler_dist_accept,
     /* .apply             = */ llama_sampler_dist_apply,
     /* .reset             = */ llama_sampler_dist_reset,
     /* .clone             = */ llama_sampler_dist_clone,
@@ -1234,6 +1392,8 @@ static struct llama_sampler_i llama_sampler_dist_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_dist_backend_apply,
     /* .backend_set_input = */ llama_sampler_dist_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_dist_backend_reset,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_dist>,
 };
 
 struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
@@ -1245,11 +1405,36 @@ struct llama_sampler * llama_sampler_init_dist(uint32_t seed) {
             /* .seed                      = */ seed,
             /* .seed_cur                  = */ seed_cur,
             /* .rng                       = */ std::mt19937(seed_cur),
-            /* .inp_uniform = */ nullptr,
+            /* .backend_transactional     = */ false,
+            /* .rng_backend               = */ std::mt19937(seed_cur),
+            /* .n_backend_draws_generated = */ 0,
+            /* .n_backend_draws_committed = */ 0,
+            /* .inp_uniforms              = */ {},
         }
     );
 }
 
+void llama_sampler_backend_begin(llama_sampler * sampler) {
+    GGML_ASSERT(sampler != nullptr);
+
+    if (sampler->iface == &llama_sampler_chain_i) {
+        auto * chain = (llama_sampler_chain *) sampler->ctx;
+        for (auto & entry : chain->samplers) {
+            if (!entry.is_backend) {
+                break;
+            }
+            llama_sampler_backend_begin(entry.ptr);
+        }
+    } else if (sampler->iface == &llama_sampler_dist_i) {
+        auto * ctx = (llama_sampler_dist *) sampler->ctx;
+        if (ctx->backend_transactional) {
+            ctx->rng_backend = ctx->rng;
+            ctx->n_backend_draws_generated = 0;
+            ctx->n_backend_draws_committed = 0;
+        }
+    }
+}
+
 // top-k
 
 struct llama_sampler_top_k : public llama_sampler_backend {
@@ -1277,8 +1462,10 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_top_k *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
@@ -1325,6 +1512,8 @@ static struct llama_sampler_i llama_sampler_top_k_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_top_k_backend_apply,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_top_k>,
 };
 
 struct llama_sampler * llama_sampler_init_top_k(int32_t k) {
@@ -1423,8 +1612,10 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_top_p *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
@@ -1521,6 +1712,8 @@ static struct llama_sampler_i llama_sampler_top_p_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_top_p_backend_apply,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_top_p>,
 };
 
 struct llama_sampler * llama_sampler_init_top_p(float p, size_t min_keep) {
@@ -1618,8 +1811,10 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_min_p *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
@@ -1680,6 +1875,8 @@ static struct llama_sampler_i llama_sampler_min_p_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_min_p_backend_apply,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_min_p>,
 };
 
 struct llama_sampler * llama_sampler_init_min_p(float p, size_t min_keep) {
@@ -1790,6 +1987,8 @@ static struct llama_sampler_i llama_sampler_typical_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_typical(float p, size_t min_keep) {
@@ -1866,8 +2065,10 @@ static void llama_sampler_backend_temp_sampling(
 
 static bool llama_sampler_temp_backend_init(
         struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_temp *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
@@ -1896,6 +2097,8 @@ static struct llama_sampler_i llama_sampler_temp_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_temp_backend_apply,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_temp>,
 };
 
 struct llama_sampler * llama_sampler_init_temp(float temp) {
@@ -2009,8 +2212,10 @@ 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) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_temp_ext *) smpl->ctx;
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     const bool res = llama_sampler_backend_support(smpl, buft);
 
@@ -2095,6 +2300,8 @@ static struct llama_sampler_i llama_sampler_temp_ext_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_temp_ext_backend_apply,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_temp_ext>,
 };
 
 struct llama_sampler * llama_sampler_init_temp_ext(float temp, float delta, float exponent) {
@@ -2202,6 +2409,8 @@ static struct llama_sampler_i llama_sampler_xtc_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_xtc(float p, float t, size_t min_keep, uint32_t seed) {
@@ -2290,7 +2499,7 @@ static struct llama_sampler * llama_sampler_mirostat_clone(const struct llama_sa
 
     // copy the state
     {
-        auto * result_ctx = (llama_sampler_mirostat *) smpl->ctx;
+        auto * result_ctx = (llama_sampler_mirostat *) result->ctx;
 
         result_ctx->mu  = ctx->mu;
         result_ctx->rng = ctx->rng;
@@ -2321,6 +2530,8 @@ static struct llama_sampler_i llama_sampler_mirostat_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_mirostat(int32_t n_vocab, uint32_t seed, float tau, float eta, int32_t m) {
@@ -2425,6 +2636,8 @@ static struct llama_sampler_i llama_sampler_mirostat_v2_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_mirostat_v2(uint32_t seed, float tau, float eta) {
@@ -2546,6 +2759,8 @@ static struct llama_sampler_i llama_sampler_grammar_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 static struct llama_sampler * llama_sampler_init_grammar_impl(
@@ -2661,6 +2876,12 @@ struct llama_sampler_penalties : public llama_sampler_backend {
     std::vector<int32_t> host_token_ids;
     std::vector<int32_t> host_counts;
 
+    void copy_state(const llama_sampler_penalties & src) {
+        // note: inp_token_ids/inp_counts belong to the current sampling graph
+        prev        = src.prev;
+        token_count = src.token_count;
+    }
+
     static bool is_disabled(
             int32_t penalty_last_n,
             float   penalty_repeat,
@@ -2790,9 +3011,15 @@ static void llama_sampler_penalties_free(struct llama_sampler * smpl) {
 
 static bool llama_sampler_penalties_backend_init(
         struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     auto * sctx = (llama_sampler_penalties *) smpl->ctx;
 
+    if (n_outputs_max_per_seq > 1) {
+        sctx->init(false);
+        return false;
+    }
+
     const bool res = llama_sampler_backend_support(smpl, buft);
 
     sctx->init(res);
@@ -2952,6 +3179,12 @@ static void llama_sampler_penalties_backend_set_input(struct llama_sampler * smp
     ggml_backend_tensor_set(sctx->inp_counts,    sctx->host_counts.data(),    0, sctx->n_max * sizeof(int32_t));
 }
 
+static void llama_sampler_penalties_backend_reset(struct llama_sampler * smpl) {
+    auto * sctx = (llama_sampler_penalties *) smpl->ctx;
+    sctx->inp_token_ids = nullptr;
+    sctx->inp_counts    = nullptr;
+}
+
 static struct llama_sampler_i llama_sampler_penalties_i = {
     /* .name              = */ llama_sampler_penalties_name,
     /* .accept            = */ llama_sampler_penalties_accept,
@@ -2963,6 +3196,8 @@ static struct llama_sampler_i llama_sampler_penalties_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_penalties_backend_apply,
     /* .backend_set_input = */ llama_sampler_penalties_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_penalties_backend_reset,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_penalties>,
 };
 
 struct llama_sampler * llama_sampler_init_penalties(
@@ -3058,6 +3293,8 @@ static struct llama_sampler_i llama_sampler_top_n_sigma_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
@@ -3395,6 +3632,8 @@ static struct llama_sampler_i llama_sampler_dry_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ 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) {
@@ -3614,6 +3853,8 @@ static struct llama_sampler_i llama_sampler_adaptive_p_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ nullptr,
     /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_adaptive_p(
@@ -3715,6 +3956,9 @@ static void llama_sampler_logit_bias_backend_apply(
 
     const size_t n = sctx->logit_bias.size();
 
+    if (sctx->inp_logit_bias == nullptr) {
+        GGML_ASSERT(sctx->inp_logit_idxs == nullptr);
+
         sctx->inp_logit_bias = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 1, n);
         ggml_set_name(sctx->inp_logit_bias, "logit_bias");
         ggml_set_input(sctx->inp_logit_bias);
@@ -3722,6 +3966,7 @@ static void llama_sampler_logit_bias_backend_apply(
         sctx->inp_logit_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, n);
         ggml_set_name(sctx->inp_logit_idxs, "logit_idxs");
         ggml_set_input(sctx->inp_logit_idxs);
+    }
 
     ggml_tensor * cur = ggml_fill(ctx, data->logits, 0.0f);
 
@@ -3756,10 +4001,18 @@ static void llama_sampler_logit_bias_backend_set_input(struct llama_sampler * sm
     ggml_backend_tensor_set(sctx->inp_logit_idxs, data_logit_idxs.data(), 0, ggml_nbytes(sctx->inp_logit_idxs));
 }
 
+static void llama_sampler_logit_bias_backend_reset(struct llama_sampler * smpl) {
+    auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
+    sctx->inp_logit_bias = nullptr;
+    sctx->inp_logit_idxs = nullptr;
+}
+
 static bool llama_sampler_logit_bias_backend_init(
         struct llama_sampler       * smpl,
-        ggml_backend_buffer_type_t   buft) {
+        ggml_backend_buffer_type_t   buft,
+        uint32_t                     n_outputs_max_per_seq) {
     GGML_UNUSED(buft);
+    GGML_UNUSED(n_outputs_max_per_seq);
 
     auto * sctx = (llama_sampler_logit_bias *) smpl->ctx;
 
@@ -3783,6 +4036,8 @@ static struct llama_sampler_i llama_sampler_logit_bias_i = {
     /* .backend_accept    = */ nullptr,
     /* .backend_apply     = */ llama_sampler_logit_bias_backend_apply,
     /* .backend_set_input = */ llama_sampler_logit_bias_backend_set_input,
+    /* .backend_reset     = */ llama_sampler_logit_bias_backend_reset,
+    /* .copy_state        = */ llama_sampler_backend_copy_state<llama_sampler_logit_bias>,
 };
 
 struct llama_sampler * llama_sampler_init_logit_bias(
@@ -4022,10 +4277,12 @@ static struct llama_sampler_i llama_sampler_infill_i = {
     /* .reset             = */ nullptr,
     /* .clone             = */ llama_sampler_infill_clone,
     /* .free              = */ llama_sampler_infill_free,
-    /* .backend_apply     = */ nullptr,
-    /* .backend_accept    = */ nullptr,
-    /* .backend_set_input = */ nullptr,
     /* .backend_init      = */ nullptr,
+    /* .backend_accept    = */ nullptr,
+    /* .backend_apply     = */ nullptr,
+    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
 };
 
 struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * vocab) {
@@ -4039,6 +4296,32 @@ struct llama_sampler * llama_sampler_init_infill(const struct llama_vocab * voca
     );
 }
 
+void llama_sampler_copy(const struct llama_sampler * src, struct llama_sampler * dst) {
+    if (!src || !dst || src == dst) {
+        return;
+    }
+
+    GGML_ASSERT(src->iface == dst->iface && "llama_sampler_copy: cannot copy between different sampler types");
+
+    if (dst->iface->copy_state) {
+        dst->iface->copy_state(src, dst);
+        return;
+    }
+
+    // build a temporary sampler carrying src's current state
+    llama_sampler * tmp = llama_sampler_clone(src);
+
+    // free dst's old state (frees dst->ctx, including children for a chain)
+    if (dst->iface->free) {
+        dst->iface->free(dst);
+    }
+
+    // transplant tmp's state into dst, then destroy the (now empty) temp shell
+    dst->ctx = tmp->ctx;
+    tmp->ctx = nullptr;
+    delete tmp;
+}
+
 // utils
 
 uint32_t llama_sampler_get_seed(const struct llama_sampler * smpl) {

src/llama-sampler.h

@@ -15,6 +15,8 @@ struct llama_sampler_chain {
     // has .backend_init() been called?
     bool is_init = false;
 
+    uint32_t n_nodes = 0;
+
     struct info {
         bool is_backend;
 
@@ -33,6 +35,9 @@ struct llama_sampler_chain {
     mutable int32_t n_sample;
 };
 
+uint32_t llama_sampler_backend_n_nodes(const llama_sampler * sampler);
+void llama_sampler_backend_begin(llama_sampler * sampler);
+
 struct llama_sampler * llama_sampler_init_dry_testing(
         float   dry_multiplier,
         float   dry_base,

tests/test-arg-parser.cpp

@@ -2,7 +2,9 @@
 #include "common.h"
 #include "download.h"
 #include "llama.h"
+#include "speculative.h"
 
+#include <limits>
 #include <string>
 #include <vector>
 #include <sstream>
@@ -14,6 +16,34 @@
 static void test(void) {
     common_params params;
 
+    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( 6, 2,  3, 6, 4);
+    assert_output_limits( 2, 1,  3, 2, 2);
+    assert_output_limits(
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max(),
+            std::numeric_limits<int32_t>::max());
+
+    {
+        common_params base;
+        base.n_parallel = 4;
+        base.n_outputs_max_per_seq = 8;
+
+        const auto draft = common_base_params_to_speculative(base);
+        assert(draft.n_outputs_max == 4);
+        assert(draft.n_outputs_max_per_seq == 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++) {
         try {

tests/test-backend-sampler.cpp

@@ -14,6 +14,7 @@
 #include <fstream>
 #include <functional>
 #include <map>
+#include <random>
 #include <string>
 #include <unordered_map>
 #include <unordered_set>
@@ -80,7 +81,13 @@ struct test_context {
     std::unordered_map<llama_seq_id, int32_t> seq_positions;
     std::unordered_map<llama_seq_id, int32_t> last_batch_info;
 
-    test_context(const test_params & params, std::vector<llama_sampler_seq_config> & configs, int32_t n_seq_max = -1) {
+    test_context(
+            const test_params & params,
+            std::vector<llama_sampler_seq_config> & configs,
+            int32_t n_seq_max = -1,
+            uint32_t n_outputs_max = 0,
+            uint32_t n_ubatch = 0,
+            uint32_t n_outputs_max_per_seq = 1) {
         auto * model = params.model.get();
 
         GGML_ASSERT(model);
@@ -89,6 +96,11 @@ struct test_context {
         llama_context_params cparams = llama_context_default_params();
         cparams.n_ctx = 512;
         cparams.n_batch = 512;
+        if (n_ubatch > 0) {
+            cparams.n_ubatch = n_ubatch;
+        }
+        cparams.n_outputs_max = n_outputs_max;
+        cparams.n_outputs_max_per_seq = n_outputs_max_per_seq;
         cparams.samplers = configs.data();
         cparams.n_samplers = configs.size();
         cparams.kv_unified = true;
@@ -262,6 +274,66 @@ struct test_context {
     }
 };
 
+struct test_single_output_backend_sampler {
+    bool backend_initialized = false;
+    uint32_t backend_outputs_max_per_seq = 0;
+    int backend_apply_count = 0;
+    int apply_count = 0;
+};
+
+static const char * test_single_output_backend_sampler_name(const llama_sampler * /*smpl*/) {
+    return "single-output-backend";
+}
+
+static void test_single_output_backend_sampler_apply(
+        llama_sampler * smpl, llama_token_data_array * /*cur_p*/) {
+    auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
+    ctx->apply_count++;
+}
+
+static void test_single_output_backend_sampler_free(llama_sampler * smpl) {
+    delete (test_single_output_backend_sampler *) smpl->ctx;
+}
+
+static bool test_single_output_backend_sampler_backend_init(
+        llama_sampler * smpl, ggml_backend_buffer_type_t /*buft*/, uint32_t n_outputs_max_per_seq) {
+    auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
+    ctx->backend_outputs_max_per_seq = n_outputs_max_per_seq;
+    if (n_outputs_max_per_seq > 1) {
+        return false;
+    }
+    ctx->backend_initialized = true;
+    return true;
+}
+
+static void test_single_output_backend_sampler_backend_apply(
+        llama_sampler * smpl, ggml_context * /*ctx*/, ggml_cgraph * /*gf*/, llama_sampler_data * /*data*/) {
+    auto * ctx = (test_single_output_backend_sampler *) smpl->ctx;
+    ctx->backend_apply_count++;
+}
+
+static llama_sampler_i test_single_output_backend_sampler_i = {
+    /* .name              = */ test_single_output_backend_sampler_name,
+    /* .accept            = */ nullptr,
+    /* .apply             = */ test_single_output_backend_sampler_apply,
+    /* .reset             = */ nullptr,
+    /* .clone             = */ nullptr,
+    /* .free              = */ test_single_output_backend_sampler_free,
+    /* .backend_init      = */ test_single_output_backend_sampler_backend_init,
+    /* .backend_accept    = */ nullptr,
+    /* .backend_apply     = */ test_single_output_backend_sampler_backend_apply,
+    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
+};
+
+static llama_sampler * test_single_output_backend_sampler_init(
+        test_single_output_backend_sampler ** sampler_ctx) {
+    auto * ctx = new test_single_output_backend_sampler;
+    *sampler_ctx = ctx;
+    return llama_sampler_init(&test_single_output_backend_sampler_i, ctx);
+}
+
 static void test_backend_greedy_sampling(const test_params & params) {
     const int seq_id = 0;
 
@@ -661,7 +733,7 @@ static void test_backend_multi_sequence_sampling(const test_params & params) {
 }
 
 static void test_backend_dist_sampling(const test_params & params) {
-    const int seq_id = 189;
+    const int seq_id = 0;
     const int32_t seed = 88;
 
     struct llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
@@ -1527,43 +1599,398 @@ static void test_backend_cpu_mixed_batch(const test_params & params) {
     printf("backend-cpu mixed batch test PASSED\n");
 }
 
-static void test_backend_max_outputs(const test_params & params) {
+static void test_backend_multi_output_limit(const test_params & params) {
-    const int seq_id = 0;
+    const llama_seq_id seq_id = 0;
-    const int32_t seed = 88;
 
-    llama_sampler_chain_params backend_chain_params = llama_sampler_chain_default_params();
+    llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
-    llama_sampler_ptr backend_sampler_chain(llama_sampler_chain_init(backend_chain_params));
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(88));
-    llama_sampler_chain_add(backend_sampler_chain.get(), llama_sampler_init_dist(seed));
+    std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
-    std::vector<llama_sampler_seq_config> backend_sampler_configs = {{ seq_id, backend_sampler_chain.get() }};
+    test_context test_ctx(params, configs, 1, 3, 0, 2);
 
-    test_context test_ctx(params, backend_sampler_configs);
+    llama_batch batch = llama_batch_init(3, 0, 1);
-
+    for (int i = 0; i < 3; ++i) {
-    llama_batch batch = llama_batch_init(512, 0, 1);
+        common_batch_add(batch, llama_vocab_bos(test_ctx.vocab), i, { seq_id }, true);
-    std::string prompt = "Hello";
-
-    std::vector<llama_token> tokens;
-    tokens.push_back(llama_vocab_bos(test_ctx.vocab));
-
-    std::vector<llama_token> prompt_tokens(32);
-    int n_tokens = llama_tokenize(test_ctx.vocab, prompt.c_str(), prompt.length(),
-                                   prompt_tokens.data(), prompt_tokens.size(),
-                                   false, false);
-    for (int i = 0; i < n_tokens; i++) {
-        tokens.push_back(prompt_tokens[i]);
     }
 
-    for (size_t i = 0; i < tokens.size(); i++) {
+    printf(">>> test_backend_multi_output_limit expected error start:\n");
-        // set all tokens as output to trigger error
-        common_batch_add(batch, tokens[i], i, { seq_id }, true);
-    }
-
-    printf(">>> test_max_outputs expected error start:\n");
     const int ret = llama_decode(test_ctx.ctx.get(), batch);
-    GGML_ASSERT(ret != 0 && "llama_decode should not succeed multiple outputs per sequence");
+    GGML_ASSERT(ret != 0 && "llama_decode should reject outputs above the per-sequence limit");
-    printf("<<< test_max_outputs expected error end.\n");
+    printf("<<< test_backend_multi_output_limit expected error end.\n");
+
     llama_batch_free(batch);
 
-    printf("backend max outputs test PASSED\n");
+    printf("backend multi-output limit 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<llama_sampler_seq_config> configs = {
+        { 0, chain_0.get() },
+        { 1, chain_1.get() },
+    };
+    test_context test_ctx(params, configs, 2, 4, 0, 2);
+
+    std::vector<llama_sampler_seq_config> reference_configs;
+    test_context reference_ctx(params, reference_configs, 2, 4);
+
+    const llama_token seq_tokens[2][2] = {
+        { llama_vocab_bos(vocab), llama_vocab_eos(vocab) },
+        { llama_vocab_eos(vocab), llama_vocab_bos(vocab) },
+    };
+
+    llama_batch batch = llama_batch_init(4, 0, 1);
+    for (int pos = 0; pos < 2; ++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<double> reference_dist(0.0, 1.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);
+
+        llama_sampler * chain = seq_id == 0 ? chain_0.get() : chain_1.get();
+        const llama_token backend_token = llama_sampler_sample(chain, 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);
+    }
+
+    llama_batch_free(batch);
+
+    printf("backend multi-sequence multi-output dist test PASSED\n");
+}
+
+static void test_backend_multi_output_dist_transaction(const test_params & params) {
+    const llama_seq_id seq_id = 0;
+    const uint32_t seed = 95;
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+
+    llama_sampler_ptr chain(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_temp(10.0f));
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
+    std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+    test_context test_ctx(params, configs, 1, 3, 2, 3);
+
+    auto verify_random = [&](int32_t row, float rnd, bool accept = true) {
+        const llama_token token = accept ?
+            llama_sampler_sample(chain.get(), test_ctx.ctx.get(), row) :
+            llama_get_sampled_token_ith(test_ctx.ctx.get(), row);
+        const float * probs = llama_get_sampled_probs_ith(test_ctx.ctx.get(), row);
+
+        GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
+        GGML_ASSERT(probs != nullptr);
+
+        float cumsum_before = 0.0f;
+        for (llama_token i = 0; i < token; ++i) {
+            cumsum_before += probs[i];
+        }
+
+        const float cumsum_sampled = cumsum_before + probs[token];
+        GGML_ASSERT(rnd >= cumsum_before - 1e-4f);
+        GGML_ASSERT(rnd <= cumsum_sampled + 1e-4f);
+    };
+
+    std::mt19937 rng(seed);
+    std::uniform_real_distribution<double> dist(0.0, 1.0);
+    float randoms[3];
+    for (float & rnd : randoms) {
+        rnd = dist(rng);
+    }
+
+    int32_t pos = 0;
+    auto decode = [&]() {
+        llama_batch batch = llama_batch_init(3, 0, 1);
+        for (int32_t i = 0; i < 3; ++i) {
+            common_batch_add(batch, llama_vocab_bos(vocab), pos++, { seq_id }, true);
+        }
+        GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+        return batch;
+    };
+
+    llama_batch batch = decode();
+    verify_random(0, randoms[0], false);
+    llama_batch_free(batch);
+
+    batch = decode();
+    verify_random(0, randoms[0]);
+    verify_random(1, randoms[1]);
+    llama_batch_free(batch);
+
+    batch = decode();
+    llama_sampler_ptr saved(llama_sampler_clone(chain.get()));
+    verify_random(0, randoms[2]);
+    llama_batch_free(batch);
+
+    llama_sampler_copy(saved.get(), chain.get());
+
+    batch = decode();
+    verify_random(0, randoms[2]);
+    llama_batch_free(batch);
+
+    printf("backend multi-output dist transaction test PASSED\n");
+}
+
+static void test_backend_multi_output_sampling_chain(const test_params & params) {
+    const llama_seq_id seq_id = 0;
+    const uint32_t seed = 88;
+    const float p = 0.9f;
+    const float temp = 0.8f;
+    const float cdf_epsilon = 1e-4f;
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+    const int32_t n_vocab = llama_vocab_n_tokens(vocab);
+    const uint32_t k = std::min<uint32_t>(512, n_vocab);
+    const llama_logit_bias bias = { llama_vocab_bos(vocab), -0.1f };
+
+    auto make_filter_chain = [&]() {
+        llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_logit_bias(n_vocab, 1, &bias));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_top_p(p, 1));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_min_p(0.01f, 1));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_temp(temp));
+        return result;
+    };
+
+    llama_sampler_ptr chain = make_filter_chain();
+    llama_sampler_chain_add(chain.get(), llama_sampler_init_dist(seed));
+    std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+    test_context test_ctx(params, configs, 1, 2, 2, 2);
+
+    std::vector<llama_sampler_seq_config> reference_configs;
+    test_context reference_ctx(params, reference_configs, 1, 2, 2);
+
+    llama_sampler_ptr reference_bias(llama_sampler_init_logit_bias(n_vocab, 1, &bias));
+    llama_sampler_ptr reference_top_k(llama_sampler_init_top_k(k));
+    llama_sampler_ptr reference_top_p(llama_sampler_init_top_p(p, 1));
+    llama_sampler_ptr reference_min_p(llama_sampler_init_min_p(0.01f, 1));
+    llama_sampler_ptr reference_temp(llama_sampler_init_temp(temp));
+    std::vector<llama_token_data> reference_data(n_vocab);
+
+    auto make_batch = [&](int32_t pos) {
+        llama_batch batch = llama_batch_init(2, 0, 1);
+        for (int i = 0; i < 2; ++i) {
+            common_batch_add(batch, llama_vocab_bos(vocab), pos + i, { seq_id }, true);
+        }
+        return batch;
+    };
+
+    llama_batch batch = make_batch(0);
+    GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+    GGML_ASSERT(llama_decode(reference_ctx.ctx.get(), batch) == 0);
+
+    for (int i = 0; i < batch.n_tokens; ++i) {
+        const llama_token backend_token = llama_sampler_sample(chain.get(), 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 llama_token * sampled_candidates = llama_get_sampled_candidates_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 uint32_t n_candidates = llama_get_sampled_candidates_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(sampled_candidates != nullptr);
+        GGML_ASSERT(reference_logits != nullptr);
+        GGML_ASSERT(n_logits == k);
+        GGML_ASSERT(n_probs == n_logits);
+        GGML_ASSERT(n_candidates == n_logits);
+
+        for (llama_token token = 0; token < n_vocab; ++token) {
+            reference_data[token] = { token, reference_logits[token], 0.0f };
+        }
+
+        llama_token_data_array reference = {
+            /* .data     = */ reference_data.data(),
+            /* .size     = */ reference_data.size(),
+            /* .selected = */ LLAMA_TOKEN_NULL,
+            /* .sorted   = */ false,
+        };
+
+        llama_sampler_apply(reference_bias.get(), &reference);
+        llama_sampler_apply(reference_top_k.get(), &reference);
+        llama_sampler_apply(reference_top_p.get(), &reference);
+        GGML_ASSERT(reference.size > 0);
+
+        float cdf = 0.0f;
+        for (size_t j = 0; j < reference.size; ++j) {
+            cdf += reference.data[j].p;
+        }
+        const float cdf_before = cdf - reference.data[reference.size - 1].p;
+        const float boundary_distance = std::min(std::fabs(cdf_before - p), std::fabs(cdf - p));
+
+        llama_sampler_apply(reference_min_p.get(), &reference);
+        llama_sampler_apply(reference_temp.get(), &reference);
+
+        std::unordered_map<llama_token, float> reference_by_id;
+        for (size_t j = 0; j < reference.size; ++j) {
+            reference_by_id.emplace(reference.data[j].id, reference.data[j].logit);
+        }
+        size_t n_backend_only = 0;
+        int32_t sampled_index = -1;
+        float prob_sum = 0.0f;
+
+        for (uint32_t j = 0; j < n_logits; ++j) {
+            GGML_ASSERT(sampled_candidates[j] >= 0 && sampled_candidates[j] < n_vocab);
+            GGML_ASSERT(std::isfinite(sampled_probs[j]));
+            GGML_ASSERT(sampled_probs[j] >= 0.0f);
+            prob_sum += sampled_probs[j];
+
+            if (sampled_candidates[j] == backend_token) {
+                sampled_index = j;
+            }
+            if (!std::isfinite(sampled_logits[j])) {
+                GGML_ASSERT(std::isinf(sampled_logits[j]) && sampled_logits[j] < 0.0f);
+                GGML_ASSERT(sampled_probs[j] == 0.0f);
+                continue;
+            }
+
+            const auto match = reference_by_id.find(sampled_candidates[j]);
+            if (match == reference_by_id.end()) {
+                ++n_backend_only;
+                continue;
+            }
+
+            const float tolerance = 1e-4f * std::max(1.0f, std::fabs(match->second));
+            GGML_ASSERT(std::fabs(sampled_logits[j] - match->second) <= tolerance);
+            reference_by_id.erase(match);
+        }
+
+        const size_t n_reference_only = reference_by_id.size();
+
+        if (n_backend_only != 0 || n_reference_only != 0) {
+            GGML_ASSERT(n_backend_only <= 1);
+            GGML_ASSERT(n_reference_only <= 1);
+            GGML_ASSERT(boundary_distance <= cdf_epsilon);
+        }
+
+        GGML_ASSERT(sampled_index >= 0);
+        GGML_ASSERT(std::isfinite(sampled_logits[sampled_index]));
+        GGML_ASSERT(sampled_probs[sampled_index] > 0.0f);
+        GGML_ASSERT(std::fabs(prob_sum - 1.0f) <= 1e-3f);
+    }
+
+    llama_batch_free(batch);
+
+    batch = make_batch(2);
+    GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+    llama_batch_free(batch);
+
+    printf("backend multi-output sampling chain test PASSED\n");
+}
+
+static void test_backend_multi_output_cpu_suffix(const test_params & params) {
+    const llama_seq_id seq_id = 0;
+    const int32_t k = 8;
+    const llama_vocab * vocab = llama_model_get_vocab(params.model.get());
+
+    auto make_chain = [&](test_single_output_backend_sampler ** sampler_ctx) {
+        llama_sampler_ptr result(llama_sampler_chain_init(llama_sampler_chain_default_params()));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_top_k(k));
+        llama_sampler_chain_add(result.get(), test_single_output_backend_sampler_init(sampler_ctx));
+        llama_sampler_chain_add(result.get(), llama_sampler_init_dist(88));
+        return result;
+    };
+
+    {
+        test_single_output_backend_sampler * sampler_ctx = nullptr;
+        llama_sampler_ptr chain = make_chain(&sampler_ctx);
+        std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+        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_outputs_max_per_seq == 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);
+
+        llama_batch_free(batch);
+    }
+
+    {
+        test_single_output_backend_sampler * sampler_ctx = nullptr;
+        llama_sampler_ptr chain = make_chain(&sampler_ctx);
+        std::vector<llama_sampler_seq_config> configs = {{ seq_id, chain.get() }};
+        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) {
+            common_batch_add(batch, llama_vocab_bos(vocab), i, { seq_id }, true);
+        }
+        GGML_ASSERT(llama_decode(test_ctx.ctx.get(), batch) == 0);
+
+        GGML_ASSERT(!sampler_ctx->backend_initialized);
+        GGML_ASSERT(sampler_ctx->backend_outputs_max_per_seq == 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);
+            GGML_ASSERT(llama_get_sampled_logits_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
+            GGML_ASSERT(llama_get_sampled_candidates_count_ith(test_ctx.ctx.get(), i) == (uint32_t) k);
+            const llama_token token = llama_sampler_sample(chain.get(), test_ctx.ctx.get(), i);
+            GGML_ASSERT(token >= 0 && token < llama_vocab_n_tokens(vocab));
+        }
+        GGML_ASSERT(sampler_ctx->apply_count == batch.n_tokens);
+
+        llama_batch_free(batch);
+    }
+
+    printf("backend multi-output CPU suffix test PASSED\n");
 }
 
 struct backend_test_case {
@@ -1583,7 +2010,11 @@ 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  },
-    { "max_outputs",     test_backend_max_outputs,             true  },
+    { "multi_output_limit",    test_backend_multi_output_limit,      true },
+    { "multi_sequence_multi_output_dist", test_backend_multi_sequence_multi_output_dist, true },
+    { "multi_output_dist_transaction", test_backend_multi_output_dist_transaction, 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  },
     { "min_p",           test_backend_min_p_sampling,          true  },
     { "cpu_mixed",       test_backend_cpu_mixed_batch,         true  },

tests/test-sampling.cpp

@@ -61,6 +61,35 @@ private:
     std::vector<llama_token_data> cur;
 };
 
+static llama_token sample_dist(llama_sampler * sampler, const std::vector<float> & logits) {
+    std::vector<llama_token_data> cur;
+    for (llama_token token_id = 0; token_id < (llama_token) logits.size(); ++token_id) {
+        cur.push_back({ token_id, logits[token_id], 0.0f });
+    }
+
+    llama_token_data_array cur_p = { cur.data(), cur.size(), -1, false };
+    llama_sampler_apply(sampler, &cur_p);
+    GGML_ASSERT(cur_p.selected >= 0);
+    GGML_ASSERT((size_t) cur_p.selected < cur_p.size);
+    return cur_p.data[cur_p.selected].id;
+}
+
+static void test_dist_singleton_rng() {
+    llama_sampler * singleton = llama_sampler_init_dist(4242);
+    llama_sampler * control   = llama_sampler_init_dist(4242);
+
+    sample_dist(singleton, { 0.0f });
+    sample_dist(control,   { 0.0f, 0.0f });
+
+    const std::vector<float> logits(256, 0.0f);
+    for (int i = 0; i < 4; ++i) {
+        GGML_ASSERT(sample_dist(singleton, logits) == sample_dist(control, logits));
+    }
+
+    llama_sampler_free(singleton);
+    llama_sampler_free(control);
+}
+
 static void test_temp(const std::vector<float> & probs, const std::vector<float> & probs_expected, float temp) {
     sampler_tester tester(probs, probs_expected);
 
@@ -308,6 +337,8 @@ static void test_perf() {
 int main(void) {
     ggml_time_init();
 
+    test_dist_singleton_rng();
+
     test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 1.0f);
     test_temp({0.1f, 0.2f, 0.3f, 0.4f}, {0.0f, 0.0f, 0.0f, 1.0f}, 0.0f);
 

tools/server/server-context.cpp

@@ -39,19 +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) {
-    const uint32_t n_batch  = params.n_batch;
-
     if (params.embedding ||
             (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
-        return n_batch;
+        return { params.n_batch, 1 };
     }
 
-    const uint32_t n_outputs_per_seq = 1 + common_speculative_n_max(&params.speculative);
+    auto result = common_speculative_get_output_limits(
+            params.n_batch, params.n_parallel, common_speculative_n_max(&params.speculative));
 
-    const uint64_t n_outputs = (uint64_t) params.n_parallel * n_outputs_per_seq;
+    result.total   = std::max<int32_t>(1, result.total);
-
+    result.per_seq = std::max<int32_t>(1, result.per_seq);
-    return std::max<uint32_t>(1, std::min<uint64_t>(n_batch, n_outputs));
+    return result;
 }
 
 // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283
@@ -1063,7 +1062,9 @@ private:
         const bool is_resume = sleeping;
 
         params_base = params;
-        params_base.n_outputs_max = server_n_outputs_max(params_base);
+        const auto output_limits = server_output_limits(params_base);
+        params_base.n_outputs_max = output_limits.total;
+        params_base.n_outputs_max_per_seq = output_limits.per_seq;
 
         const bool has_mmproj = !params.mmproj.path.empty();
         const bool has_draft = params.speculative.has_dft();
@@ -1832,18 +1833,13 @@ private:
 
             const bool need_pre_sample_logits = task.params.sampling.n_probs > 0 && !task.params.post_sampling_probs;
 
-            bool backend_sampling = true;
+            bool use_backend_sampling = task.params.sampling.backend_sampling;
-
-            backend_sampling &= task.params.sampling.backend_sampling;
-
-            // TODO: speculative decoding requires multiple samples per batch - not supported yet
-            backend_sampling &= !(slot.can_speculate());
 
             // TODO: getting pre sampling logits is not yet supported with backend sampling
-            backend_sampling &= !need_pre_sample_logits;
+            use_backend_sampling &= !need_pre_sample_logits;
 
             // TODO: tmp until backend sampling is fully implemented
-            if (backend_sampling) {
+            if (use_backend_sampling) {
                 llama_set_sampler(ctx_tgt, slot.id, common_sampler_get(slot.smpl.get()));
             } else {
                 llama_set_sampler(ctx_tgt, slot.id, nullptr);
@@ -3865,7 +3861,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;
             }
 
@@ -3876,7 +3873,6 @@ 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_i_batch.size() == n_draft + 1);
@@ -3915,7 +3911,7 @@ private:
                         slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1);
 
                         slot.prompt.tokens.keep_first(ckpt.n_tokens);
-                        slot.smpl = std::move(smpl_save);
+                        common_sampler_copy(smpl_save.get(), slot.smpl.get());
 
                         return;
                     }

tools/server/tests/unit/test_speculative.py

@@ -25,33 +25,34 @@ def fixture_create_server():
 
 def test_with_and_without_draft():
     global server
+    request = {
+        "prompt": "I believe the meaning of life is",
+        "temperature": 0.8,
+        "top_k": 40,
+        "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={
+    res = server.make_request("POST", "/completion", data=request)
-        "prompt": "I believe the meaning of life is",
-        "temperature": 0.0,
-        "top_k": 1,
-        "n_predict": 16,
-    })
     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={
+    res = server.make_request("POST", "/completion", data=request)
-        "prompt": "I believe the meaning of life is",
-        "temperature": 0.0,
-        "top_k": 1,
-        "n_predict": 16,
-    })
     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():