From 77757a34672a888845a09ab45728b8d34aa475ff Mon Sep 17 00:00:00 2001
From: Daniel Bevenius <daniel.bevenius@gmail.com>
Date: Mon, 12 Oct 2026 12:02:24 +0200
Subject: [PATCH] examples : fix draft sampling and rejection handling (#30095)

* examples : fix draft sampling and rejection handling

This commit fixes two issues in the speculative example which are
described in detail in the linked issue.

Resolves: https://github.com/ggml-org/llama.cpp/issues/29975

* examples : add kv-unified requirement for parallel > 1
---
 common/arg.cpp                       |  2 +-
 examples/speculative/speculative.cpp | 29 ++++++++++++++++++----------
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 7c1b81e0f..6bc4a84e2 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -1730,7 +1730,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
         [](common_params & params, bool value) {
             params.kv_unified = value;
         }
-    ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL}));
+    ).set_env("LLAMA_ARG_KV_UNIFIED").set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_PERPLEXITY, LLAMA_EXAMPLE_BATCHED, LLAMA_EXAMPLE_BENCH, LLAMA_EXAMPLE_PARALLEL, LLAMA_EXAMPLE_SPECULATIVE}));
     add_opt(common_arg(
         {"--cache-idle-slots"},
         {"--no-cache-idle-slots"},
diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp
index 1bb47594b..8f347b0d4 100644
--- a/examples/speculative/speculative.cpp
+++ b/examples/speculative/speculative.cpp
@@ -50,6 +50,11 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    if (params.n_parallel > 1 && !params.kv_unified) {
+        LOG_ERR("%s: --kv-unified/-kvu is required when --parallel/-np > 1\n", __func__);
+        return 1;
+    }
+
     if (params.speculative.draft.mparams.path.empty()) {
         LOG_ERR("%s: --model-draft is required\n", __func__);
         return 1;
@@ -267,7 +272,7 @@ int main(int argc, char ** argv) {
             // for stochastic sampling, attempt to match the token with the drafted tokens
             {
                 bool accept = false;
-                if (params.sampling.temp > 0) {
+                if (params.sampling.temp > 0 && n_seq_dft == 1) {
                     // stochastic verification
                     common_sampler_sample(smpl, ctx_tgt, drafts[s_keep].i_batch_tgt[i_dft], true);
 
@@ -341,13 +346,17 @@ int main(int argc, char ** argv) {
 
                             float sum_probs = 0.0f;
 
-                            for (size_t i = 0; i < dist_tgt.size; i++) {
-                                if (i < dist_dft.size) {
-                                    dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p - dist_dft.data[i].p);
-                                } else {
-                                    dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p);
+                            for (size_t i = 0, j = 0; i < dist_tgt.size; i++) {
+                                // try to find the current target token id in the draft distribution
+                                while (j < dist_dft.size && dist_dft.data[j].id < dist_tgt.data[i].id) {
+                                    ++j;
                                 }
 
+                                // tokens absent from the draft distribution have zero probability
+                                float p_dft = j < dist_dft.size && dist_dft.data[j].id == dist_tgt.data[i].id ?
+                                    dist_dft.data[j].p : 0.0f;
+
+                                dist_tgt.data[i].p = std::max(0.0f, dist_tgt.data[i].p - p_dft);
                                 sum_probs += dist_tgt.data[i].p;
                             }
 
@@ -370,7 +379,7 @@ int main(int argc, char ** argv) {
                                 // synchronize active status for sequences with the same drafted token
                                 drafts[i].active = drafts[i].active && accept;
                                 if (!drafts[i].active) {
-                                    active_seqs.erase(s);
+                                    active_seqs.erase(i);
                                 }
                             }
                         }
@@ -394,7 +403,7 @@ int main(int argc, char ** argv) {
                         token_str = common_token_to_piece(ctx_tgt, token_id);
                     }
                 } else {
-                    // greedy verification
+                    // verify drafts by matching tokens sampled from the target
 
                     // sample from the target model
                     LOG_DBG("sampling target: s_keep = %3d, i_dft = %3d, i_batch_tgt = %3d\n", s_keep, i_dft, drafts[s_keep].i_batch_tgt[i_dft]);
@@ -518,7 +527,7 @@ int main(int argc, char ** argv) {
                     continue;
                 }
 
-                common_sampler_sample(drafts[s].smpl, ctx_dft, drafts[s].i_batch_dft, true);
+                const llama_token sampled_token = common_sampler_sample(drafts[s].smpl, ctx_dft, drafts[s].i_batch_dft, true);
 
                 const auto * cur_p = common_sampler_get_candidates(drafts[s].smpl, true);
 
@@ -567,7 +576,7 @@ int main(int argc, char ** argv) {
 
                 // add drafted token for each sequence
                 for (int is = 0; is < (int) sa.size(); ++is) {
-                    const llama_token id = cur_p->data[is].id;
+                    const llama_token id = n_seq_dft == 1 ? sampled_token : cur_p->data[is].id;
 
                     const int s = sa[is];