Commit 1ab09d632 for llama.cpp
commit 1ab09d632ebc30f171f58547a1110f8d6e9546c0
Author: Pranesh Gonegandla <pranesh.iitp@gmail.com>
Date: Fri Oct 23 05:45:36 2026 +0000
Extend probabilistic sampling for DFlash2 drafts (#30663)
* spec : add probabilistic sampling for DFlash2 drafts
* spec : apply review suggestions
---------
Co-authored-by: Pranesh Gonegandla <pgonegandla@nvidia.com>
diff --git a/common/speculative.cpp b/common/speculative.cpp
index 5a0b5b512..4305ac6ac 100644
--- a/common/speculative.cpp
+++ b/common/speculative.cpp
@@ -19,6 +19,7 @@
#include <fstream>
#include <iomanip>
#include <map>
+#include <random>
#include <cinttypes>
#define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
@@ -1010,6 +1011,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
bool is_mrope = false;
int32_t selector_top_k = 0;
+ // DFlash2 probabilistic drafting: one generator per sequence and the seed it was last seeded with (-1: reseed)
+ std::vector<std::mt19937> dflash2_rng;
+ std::vector<int64_t> dflash2_seed;
+
// draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
const bool is_dspark;
@@ -1060,6 +1065,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
selector_top_k = llama_model_dflash_selector_top_k(model_dft);
is_dflash2 = selector_top_k > 0;
+ dflash2_rng.resize(n_seq);
+ dflash2_seed.assign(n_seq, -1);
char lilicorr_meta[32] = {};
lilicorr_top_k = llama_model_meta_val_str(model_dft, "lilicorr.candidate_topk", lilicorr_meta, sizeof(lilicorr_meta)) >= 0
? std::atoi(lilicorr_meta) : 0;
@@ -1158,6 +1165,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
return;
}
+ // reseed per request rather than per round, or two identical requests differ
+ dflash2_seed[seq_id] = -1;
+
const int32_t N = (int32_t) prompt.size();
if (N <= 0) {
return;
@@ -1338,6 +1348,21 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const float * lattice = llama_get_embeddings_nextn(ctx_dft);
GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
+ // DFlash2 can instead sample each step and hand q to the verifier for rejection sampling
+ const bool sample = is_dflash2 && params.probabilistic && dp.result_q;
+
+ // greedy keeps temperature 1, sampling uses the target's
+ const float temp = sample ? dp.temp : 1.0f;
+
+ if (sample && dflash2_seed[seq_id] != dp.seed) {
+ dflash2_seed[seq_id] = dp.seed;
+ // the default seed means random; mixed so it differs from the target's, as in spec_retune
+ dflash2_rng[seq_id].seed(dp.seed == LLAMA_DEFAULT_SEED ? std::random_device{}() : dp.seed ^ 0x85ebca6bu);
+ }
+
+ std::vector<float> probs;
+ std::vector<int32_t> sorted_idx;
+
int32_t predecessor = 0;
for (int32_t i = 1; i < (is_dflash2 ? n_block_tokens : params.n_max + 1); ++i) {
const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
@@ -1346,21 +1371,53 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
predecessor = (int32_t) std::distance(scores,
std::max_element(scores, scores + top_k));
- if (params.p_min > 0.0f) {
- // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max))
+ if (params.p_min > 0.0f || sample) {
+ // softmax(scores / temp) at the argmax, i.e. 1 / sum(exp((s_k - s_max) / temp))
+ probs.resize(top_k);
float sum = 0.0f;
for (int32_t k = 0; k < top_k; ++k) {
- sum += std::exp(scores[k] - scores[predecessor]);
+ probs[k] = std::exp((scores[k] - scores[predecessor]) / temp);
+ sum += probs[k];
}
if (1.0f / sum < params.p_min) {
break;
}
+ for (int32_t k = 0; k < top_k; ++k) {
+ probs[k] /= sum;
+ }
+ }
+ if (sample) {
+ // ggml_top_k leaves candidates unordered, so walk them by token id to keep a seed's draft stable
+ sorted_idx.resize(top_k);
+ for (int32_t k = 0; k < top_k; ++k) {
+ sorted_idx[k] = k;
+ }
+ std::sort(sorted_idx.begin(), sorted_idx.end(), [row](int32_t a, int32_t b) { return row[a] < row[b]; });
+
+ // a round-off overshoot keeps the argmax, which has q > 0
+ double u = std::uniform_real_distribution<double>(0.0, 1.0)(dflash2_rng[seq_id]);
+ for (const int32_t k : sorted_idx) {
+ u -= probs[k];
+ if (u < 0.0) {
+ predecessor = k;
+ break;
+ }
+ }
+
+ auto & q = dp.result_q->emplace_back();
+ q.reserve(top_k);
+ for (int32_t k = 0; k < top_k; ++k) {
+ q.push_back({ (llama_token) row[k], scores[k], probs[k] });
+ }
}
result.push_back((llama_token) row[predecessor]);
}
if (result.size() < (size_t) params.n_min) {
result.clear();
+ if (sample) {
+ dp.result_q->clear();
+ }
}
continue;
}
diff --git a/common/speculative.h b/common/speculative.h
index a59846829..4262327cb 100644
--- a/common/speculative.h
+++ b/common/speculative.h
@@ -73,7 +73,7 @@ struct common_speculative_draft_params {
// the generated draft from the last _draft() call
llama_tokens * result;
- // candidate distribution per drafted token; set it to make draft-simple and draft-mtp sample
+ // candidate distribution per drafted token; set it to make draft-simple, draft-mtp and draft-dflash (DFlash2 only) sample
std::vector<std::vector<llama_token_data>> * result_q = nullptr;
// the target's temp and seed, read only when the drafter samples probabilistically