@miller-tech/uap 1.217.0 → 1.218.0

package/config/llama-patches/dflash2-pr27342.patch

@@ -0,0 +1,1365 @@
+diff --git a/common/sampling.cpp b/common/sampling.cpp
+index 06dea1e1c..da8bdc1b1 100644
+--- a/common/sampling.cpp
++++ b/common/sampling.cpp
+@@ -4,6 +4,7 @@
+ #include "fit.h"
+ #include "log.h"
+ #include "reasoning-budget.h"
++#include "speculative.h"
+ 
+ #include "ggml.h"
+ 
+@@ -12,6 +13,7 @@
+ #include <climits>
+ #include <cmath>
+ #include <cstring>
++#include <random>
+ #include <unordered_map>
+ #include <vector>
+ 
+@@ -121,10 +123,14 @@ struct common_sampler {
+ 
+     llama_token_data_array cur_p;
+ 
++    uint32_t speculative_seed;
++    std::mt19937 speculative_rng;
++
+     void reset() {
+         prev.clear();
+ 
+         llama_sampler_reset(chain);
++        speculative_rng.seed(speculative_seed);
+     }
+ 
+     void set_logits(struct llama_context * ctx, int idx) {
+@@ -424,6 +430,8 @@ struct common_sampler * common_sampler_init(
+         params.backend_sampling = false;
+     }
+ 
++    // Keep verifier randomness independent from both target and draft sampling.
++    const uint32_t speculative_seed = llama_sampler_get_seed(chain) ^ 0x9e3779b9U;
+     auto * result = new common_sampler {
+         /* .params  = */ params,
+         /* .grmr    = */ grmr,
+@@ -432,6 +440,8 @@ struct common_sampler * common_sampler_init(
+         /* .prev    = */ ring_buffer<llama_token>(std::max(32, params.n_prev)),
+         /* .cur     = */ {},
+         /* .cur_p   = */ {},
++        /* .speculative_seed = */ speculative_seed,
++        /* .speculative_rng  = */ std::mt19937(speculative_seed),
+     };
+ 
+     return result;
+@@ -515,6 +525,8 @@ struct common_sampler * common_sampler_clone(common_sampler * gsmpl) {
+         /* .prev    = */ gsmpl->prev,
+         /* .cur     = */ gsmpl->cur,
+         /* .cur_p   = */ gsmpl->cur_p,
++        /* .speculative_seed = */ gsmpl->speculative_seed,
++        /* .speculative_rng  = */ gsmpl->speculative_rng,
+     };
+ }
+ 
+@@ -535,6 +547,8 @@ void common_sampler_copy(const common_sampler * src, common_sampler * dst) {
+     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->speculative_seed = src->speculative_seed;
++    dst->speculative_rng  = src->speculative_rng;
+     dst->t_total_us = src->t_total_us;
+ }
+ 
+@@ -705,6 +719,79 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
+     return result;
+ }
+ 
++std::vector<llama_token> common_sampler_sample_and_accept_n(
++        struct common_sampler * gsmpl,
++        struct llama_context * ctx,
++        const std::vector<int> & idxs,
++        const llama_tokens & draft,
++        const std::vector<common_speculative_token_dist> & dists,
++        bool grammar_first) {
++    GGML_ASSERT(idxs.size() == draft.size() + 1);
++    GGML_ASSERT(dists.size() == draft.size());
++
++    std::vector<llama_token> result;
++    result.reserve(idxs.size());
++
++    std::uniform_real_distribution<float> uniform(0.0f, 1.0f);
++    size_t i = 0;
++    for (; i < draft.size(); ++i) {
++        // Residual sampling needs the target distribution after every constraint.
++        const llama_token fallback = common_sampler_sample(gsmpl, ctx, idxs[i], true);
++        const auto & q = dists[i];
++        GGML_ASSERT(q.ids.size() == q.probs.size());
++
++        std::unordered_map<llama_token, float> q_probs;
++        q_probs.reserve(q.ids.size());
++        for (size_t j = 0; j < q.ids.size(); ++j) {
++            q_probs[q.ids[j]] += q.probs[j];
++        }
++        const auto q_prob = [&](llama_token id) {
++            const auto it = q_probs.find(id);
++            return it == q_probs.end() ? 0.0f : it->second;
++        };
++
++        auto * p = common_sampler_get_candidates(gsmpl, false);
++        float p_draft = 0.0f;
++        const float q_draft = q_prob(draft[i]);
++        for (size_t j = 0; j < p->size; ++j) {
++            if (p->data[j].id == draft[i]) {
++                p_draft = p->data[j].p;
++                break;
++            }
++        }
++
++        if (q_draft > 0.0f && uniform(gsmpl->speculative_rng) * q_draft <= p_draft) {
++            common_sampler_accept(gsmpl, draft[i], true);
++            result.push_back(draft[i]);
++            continue;
++        }
++
++        std::vector<float> residual(p->size);
++        float residual_sum = 0.0f;
++        for (size_t j = 0; j < p->size; ++j) {
++            residual[j] = std::max(0.0f, p->data[j].p - q_prob(p->data[j].id));
++            residual_sum += residual[j];
++        }
++
++        llama_token id = fallback;
++        if (residual_sum > 0.0f) {
++            std::discrete_distribution<size_t> sample(residual.begin(), residual.end());
++            id = p->data[sample(gsmpl->speculative_rng)].id;
++        }
++        common_sampler_accept(gsmpl, id, true);
++        result.push_back(id);
++        break;
++    }
++
++    if (i == draft.size()) {
++        const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first);
++        common_sampler_accept(gsmpl, id, true);
++        result.push_back(id);
++    }
++
++    return result;
++}
++
+ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first) {
+     std::vector<int> idxs(draft.size() + 1);
+     for (size_t i = 0; i < idxs.size(); ++i) {
+@@ -714,6 +801,19 @@ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sample
+     return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, grammar_first);
+ }
+ 
++std::vector<llama_token> common_sampler_sample_and_accept_n(
++        struct common_sampler * gsmpl,
++        struct llama_context * ctx,
++        const llama_tokens & draft,
++        const std::vector<common_speculative_token_dist> & dists,
++        bool grammar_first) {
++    std::vector<int> idxs(draft.size() + 1);
++    for (size_t i = 0; i < idxs.size(); ++i) {
++        idxs[i] = i;
++    }
++    return common_sampler_sample_and_accept_n(gsmpl, ctx, idxs, draft, dists, grammar_first);
++}
++
+ uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl) {
+     return llama_sampler_get_seed(gsmpl->chain);
+ }
+diff --git a/common/sampling.h b/common/sampling.h
+index ced3c8364..5f3273eae 100644
+--- a/common/sampling.h
++++ b/common/sampling.h
+@@ -33,6 +33,7 @@
+ //
+ 
+ struct common_sampler;
++struct common_speculative_token_dist;
+ 
+ // llama_sampler API overloads
+ 
+@@ -85,6 +86,23 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co
+ //
+ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector<int> & idxs, const llama_tokens & draft, bool grammar_first = false);
+ 
++// maximal-coupling verification for stochastic speculative decoding
++std::vector<llama_token> common_sampler_sample_and_accept_n(
++        struct common_sampler * gsmpl,
++        struct llama_context * ctx,
++        const std::vector<int> & idxs,
++        const llama_tokens & draft,
++        const std::vector<common_speculative_token_dist> & dists,
++        bool grammar_first = false);
++
++// assume idxs == [ 0, 1, 2, ..., draft.size() ]
++std::vector<llama_token> common_sampler_sample_and_accept_n(
++        struct common_sampler * gsmpl,
++        struct llama_context * ctx,
++        const llama_tokens & draft,
++        const std::vector<common_speculative_token_dist> & dists,
++        bool grammar_first = false);
++
+ // assume idxs == [ 0, 1, 2, ..., draft.size() ]
+ std::vector<llama_token> common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false);
+ 
+diff --git a/common/speculative.cpp b/common/speculative.cpp
+index 89e9b2782..78bb4a6d5 100644
+--- a/common/speculative.cpp
++++ b/common/speculative.cpp
+@@ -14,10 +14,12 @@
+ 
+ #include <algorithm>
+ #include <cassert>
++#include <chrono>
++#include <cinttypes>
+ #include <cstring>
+ #include <iomanip>
+ #include <map>
+-#include <cinttypes>
++#include <random>
+ 
+ #define SPC_DBG(fmt, ...) LOG_DBG("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+ #define SPC_TRC(fmt, ...) LOG_TRC("spec %12.*s: " fmt, 12, __func__, __VA_ARGS__)
+@@ -923,6 +925,11 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+     int32_t     block_size    = 0;
+     llama_token mask_token_id = 0;
+ 
++    bool    is_dflash2     = false;
++    int32_t selector_top_k = 0;
++    std::vector<std::mt19937> selector_rng;
++    std::vector<bool> selector_reset;
++
+     // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
+     const bool is_dspark;
+ 
+@@ -966,6 +973,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+             if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) {
+                 sample_from_anchor = std::strcmp(buf, "true") == 0;
+             }
++            if (llama_model_meta_val_str(model_dft, "dflash.selector_top_k", buf, sizeof(buf)) >= 0) {
++                selector_top_k = std::atoi(buf);
++                is_dflash2 = selector_top_k > 0;
++            }
+         }
+         mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
+ 
+@@ -991,14 +1002,17 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+         for (auto & s : smpls) {
+             common_params_sampling sparams;
+             sparams.no_perf  = false;
+-            sparams.top_k    = 10;
++            sparams.top_k    = is_dflash2 ? selector_top_k : 10;
+             sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K };
+             s.reset(common_sampler_init(model_dft, sparams));
+         }
+ 
++        selector_rng.resize(n_seq);
++        selector_reset.assign(n_seq, true);
++
+         // offload draft sampling to the backend
+         backend_chains.assign(n_seq, nullptr);
+-        if (this->params.backend_sampling) {
++        if (this->params.backend_sampling && !is_dflash2) {
+             for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+                 llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params());
+                 llama_sampler_chain_add(chain, llama_sampler_init_top_k(10));
+@@ -1017,7 +1031,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+             llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true);
+         }
+ 
+-        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true);
++        // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits.
++        llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2);
+         llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention
+     }
+ 
+@@ -1048,6 +1063,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+             return;
+         }
+ 
++        selector_reset[seq_id] = true;
++
+         const llama_pos pos_max = llama_memory_seq_pos_max(llama_get_memory(params.ctx_dft), seq_id);
+         if (pos_max < N - 1) {
+             LOG_WRN("%s: ctx_dft pos_max=%d < N-1=%d - process() did not run on every prefill ubatch. "
+@@ -1074,87 +1091,68 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+ 
+         const int32_t n_tokens = batch_in.n_tokens;
+ 
+-        // per-seq inclusive batch range (assumes each seq's tokens are contiguous in the batch)
+-        std::vector<int32_t> i_batch_beg(n_seq, -1);
+-        std::vector<int32_t> i_batch_end(n_seq, -1);
+-        for (int32_t k = 0; k < n_tokens; ++k) {
+-            GGML_ASSERT(batch_in.n_seq_id[k] == 1);
+-            const llama_seq_id seq_id = batch_in.seq_id[k][0];
+-            if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+-                continue;
+-            }
+-            i_batch_end[seq_id] = k;
+-            if (i_batch_beg[seq_id] < 0) {
+-                i_batch_beg[seq_id] = k;
+-            }
+-        }
+-
+         auto * ctx_tgt = this->params.ctx_tgt;
+         auto * ctx_dft = this->params.ctx_dft;
+ 
+         const int32_t n_ubatch = (int32_t) llama_n_ubatch(ctx_dft);
+ 
+-        for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) {
+-            if (i_batch_beg[seq_id] < 0) {
+-                continue;
+-            }
+-            const int32_t n_rows = i_batch_end[seq_id] - i_batch_beg[seq_id] + 1;
+-
+-            for (int32_t offset = 0; offset < n_rows; offset += n_ubatch) {
+-                const int32_t n_chunk = std::min(n_ubatch, n_rows - offset);
+-
+-                // gather this chunk's target features, interleaved by extract layer
+-                features_buf.resize((size_t) n_chunk * n_embd_enc);
+-                for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
+-                    const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
+-                    if (!layer) {
+-                        GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
+-                    }
+-                    for (int32_t i = 0; i < n_chunk; ++i) {
+-                        float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
+-                        const float * src = layer + (size_t) (i_batch_beg[seq_id] + offset + i) * n_embd_tgt;
+-                        std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
+-                    }
++        // Flatten token-wise encoder work into shared chunks while preserving each row's position and sequence.
++        for (int32_t offset = 0; offset < n_tokens; offset += n_ubatch) {
++            const int32_t n_chunk = std::min(n_ubatch, n_tokens - offset);
++            features_buf.resize((size_t) n_chunk * n_embd_enc);
++            for (uint32_t k = 0; k < target_layer_ids_n; ++k) {
++                const float * layer = llama_get_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k]);
++                if (!layer) {
++                    GGML_ABORT("DFlash: target layer %d input not extracted.", target_layer_ids[k]);
++                }
++                for (int32_t i = 0; i < n_chunk; ++i) {
++                    float       * dst = features_buf.data() + (size_t) i * n_embd_enc + k * (size_t) n_embd_tgt;
++                    const float * src = layer + (size_t) (offset + i) * n_embd_tgt;
++                    std::memcpy(dst, src, (size_t) n_embd_tgt * sizeof(float));
+                 }
++            }
+ 
+-                // fuse extracted features through DFlash encoder
+-                llama_batch enc_batch = {
+-                    /*.n_tokens =*/ n_chunk,
+-                    /*.token    =*/ nullptr,
+-                    /*.embd     =*/ features_buf.data(),
+-                    /*.pos      =*/ nullptr,
+-                    /*.n_seq_id =*/ nullptr,
+-                    /*.seq_id   =*/ nullptr,
+-                    /*.logits   =*/ nullptr,
+-                };
++            llama_batch enc_batch = {
++                /*.n_tokens =*/ n_chunk,
++                /*.token    =*/ nullptr,
++                /*.embd     =*/ features_buf.data(),
++                /*.pos      =*/ nullptr,
++                /*.n_seq_id =*/ nullptr,
++                /*.seq_id   =*/ nullptr,
++                /*.logits   =*/ nullptr,
++            };
+ 
+-                int32_t rc = llama_encode(ctx_dft, enc_batch);
+-                if (rc != 0) {
+-                    LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+-                            __func__, rc, (int) n_chunk, (int) offset);
+-                    return false;
+-                }
++            int32_t rc = llama_encode(ctx_dft, enc_batch);
++            if (rc != 0) {
++                LOG_ERR("%s: llama_encode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
++                        __func__, rc, (int) n_chunk, (int) offset);
++                return false;
++            }
+ 
+-                const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
+-                GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
++            const float * inp_g = llama_get_embeddings_nextn(ctx_dft);
++            GGML_ASSERT(inp_g && "DFlash encoder produced no output.");
+ 
+-                // inject the DFlash decoder K/V cache at the tokens' target positions
+-                batch_inject.n_tokens = n_chunk;
+-                std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
++            batch_inject.n_tokens = n_chunk;
++            std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float));
++            for (int32_t i = 0; i < n_chunk; ++i) {
++                const int32_t j = offset + i;
++                GGML_ASSERT(batch_in.n_seq_id[j] == 1);
++                const llama_seq_id seq_id = batch_in.seq_id[j][0];
++                GGML_ASSERT(seq_id >= 0 && seq_id < (llama_seq_id) n_seq);
++                batch_inject.pos[i]       = batch_in.pos[j];
++                batch_inject.n_seq_id[i]  = 1;
++                batch_inject.seq_id[i][0] = seq_id;
++                batch_inject.logits[i]    = false;
++            }
+ 
+-                for (int32_t i = 0; i < n_chunk; ++i) {
+-                    batch_inject.pos[i]       = batch_in.pos[i_batch_beg[seq_id] + offset + i];
+-                    batch_inject.n_seq_id[i]  = 1;
+-                    batch_inject.seq_id[i][0] = seq_id;
+-                    batch_inject.logits[i]    = false;
+-                }
+-                rc = llama_decode(ctx_dft, batch_inject);
+-                if (rc != 0) {
+-                    LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
+-                            __func__, rc, (int) n_chunk, (int) offset);
+-                    return false;
+-                }
++            rc = llama_decode(ctx_dft, batch_inject);
++            if (rc != 0) {
++                LOG_ERR("%s: llama_decode(ctx_dft) failed rc=%d (n_tokens=%d, offset=%d)\n",
++                        __func__, rc, (int) n_chunk, (int) offset);
++                return false;
+             }
++            // The server may switch contexts before the next draft decode.
++            llama_synchronize(ctx_dft);
+         }
+ 
+         return true;
+@@ -1186,7 +1184,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+             i_block_beg[seq_id] = batch.n_tokens;
+             n_block    [seq_id] = n_block_tokens;
+             for (int32_t i = 0; i < n_block_tokens; ++i) {
+-                common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true);
++                common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2);
+             }
+         }
+ 
+@@ -1214,6 +1212,63 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
+ 
+             auto & result = *dp.result;
+ 
++            if (dp.dists) {
++                dp.dists->clear();
++            }
++
++            if (is_dflash2) {
++                GGML_ASSERT(dp.temperature <= 0.0f || dp.dists);
++                const float * lattice = llama_get_embeddings_nextn(ctx_dft);
++                GGML_ASSERT(lattice && "DFlash2 selector produced no lattice");
++
++                if (selector_reset[seq_id]) {
++                    uint32_t seed = dp.seed;
++                    if (seed == LLAMA_DEFAULT_SEED) {
++                        seed = (uint32_t) std::chrono::high_resolution_clock::now().time_since_epoch().count();
++                    }
++                    selector_rng[seq_id].seed(seed ^ 0x85ebca6bU);
++                    selector_reset[seq_id] = false;
++                }
++
++                int32_t predecessor = 0;
++                for (int32_t i = 1; i < n_block_tokens; ++i) {
++                    const float * row = lattice + (size_t) (beg + i) * n_embd_dec;
++                    const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k;
++
++                    if (dp.temperature > 0.0f) {
++                        common_speculative_token_dist dist;
++                        dist.ids.resize(selector_top_k);
++                        dist.probs.resize(selector_top_k);
++                        const float max_score = *std::max_element(scores, scores + selector_top_k);
++                        float sum = 0.0f;
++                        for (int32_t k = 0; k < selector_top_k; ++k) {
++                            dist.ids[k] = (llama_token) row[k];
++                            dist.probs[k] = std::exp((scores[k] - max_score) / dp.temperature);
++                            sum += dist.probs[k];
++                        }
++                        for (float & p : dist.probs) {
++                            p /= sum;
++                        }
++                        std::discrete_distribution<int32_t> sample(dist.probs.begin(), dist.probs.end());
++                        predecessor = sample(selector_rng[seq_id]);
++                        result.push_back(dist.ids[predecessor]);
++                        dp.dists->push_back(std::move(dist));
++                    } else {
++                        predecessor = (int32_t) std::distance(scores,
++                                std::max_element(scores, scores + selector_top_k));
++                        result.push_back((llama_token) row[predecessor]);
++                    }
++                }
++
++                if (result.size() < (size_t) params.n_min) {
++                    result.clear();
++                    if (dp.dists) {
++                        dp.dists->clear();
++                    }
++                }
++                continue;
++            }
++
+             if (is_dspark) {
+                 // DSpark: read from the first draft slot, truncate below the confidence threshold
+                 const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
+@@ -2339,7 +2394,7 @@ common_params common_base_params_to_speculative(const common_params & params) {
+     result.n_outputs_max = params.n_parallel;
+     result.n_outputs_max_per_seq = 1;
+ 
+-    // dflash/dspark decode the whole noise block in a single pass and sample every block position on the backend
++    // dflash/dspark decode every sequence's full noise block in one pass
+     // TODO: refactor such properties to be announced by the speculative types
+     //       something like `struct common_speculative_type_props common_speculative_type_get_props(...);`
+     const bool has_block_draft = std::any_of(
+@@ -2351,6 +2406,8 @@ common_params common_base_params_to_speculative(const common_params & params) {
+         // per-seq output positions: DFlash decodes anchor + n_max masks (n_max + 1); DSpark n_max -> +1 covers both
+         const int32_t per_seq = std::max(1, params_spec.n_max + 1);
+         result.n_outputs_max = params.n_parallel * per_seq;
++        result.n_batch  = std::max(result.n_batch,  result.n_outputs_max);
++        result.n_ubatch = std::max(result.n_ubatch, result.n_outputs_max);
+         if (params_spec.backend_sampling) {
+             result.n_outputs_max_per_seq = per_seq;
+         }
+diff --git a/common/speculative.h b/common/speculative.h
+index 12ae31b7d..f7dd65267 100644
+--- a/common/speculative.h
++++ b/common/speculative.h
+@@ -5,6 +5,11 @@
+ 
+ struct common_speculative;
+ 
++struct common_speculative_token_dist {
++    llama_tokens ids;
++    std::vector<float> probs;
++};
++
+ // comma separated list the provided types
+ std::string common_speculative_type_name_str(const std::vector<enum common_speculative_type> & types);
+ 
+@@ -60,6 +65,12 @@ struct common_speculative_draft_params {
+ 
+     // the generated draft from the last _draft() call
+     llama_tokens * result;
++
++    // optional sparse proposal distributions, one per draft token
++    std::vector<common_speculative_token_dist> * dists = nullptr;
++
++    float temperature = 0.0f;
++    uint32_t seed = LLAMA_DEFAULT_SEED;
+ };
+ 
+ common_speculative_draft_params & common_speculative_get_draft_params(common_speculative * spec, llama_seq_id seq_id);
+diff --git a/conversion/__init__.py b/conversion/__init__.py
+index 5ae6ad819..5d303b4a3 100644
+--- a/conversion/__init__.py
++++ b/conversion/__init__.py
+@@ -54,6 +54,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
+     "DeepseekV3ForCausalLM": "deepseek",
+     "DeepseekV32ForCausalLM": "deepseek",
+     "DFlashDraftModel": "qwen",
++    "DFlash2DraftModel": "qwen",
+     "Qwen3DSparkModel": "qwen",
+     "DSparkDraftModel": "qwen",
+     "DSparkSpeculator": "qwen",
+diff --git a/conversion/qwen.py b/conversion/qwen.py
+index 26b10452b..ab2a0fb4d 100644
+--- a/conversion/qwen.py
++++ b/conversion/qwen.py
+@@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase):
+     model_arch = gguf.MODEL_ARCH.QWEN35MOE
+ 
+ 
+-@ModelBase.register("DFlashDraftModel")
++@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel")
+ @ModelBase.example("z-lab/Qwen3.5-9B-DFlash")
+ class DFlashModel(Qwen3Model):
+     model_arch = gguf.MODEL_ARCH.DFLASH
+@@ -678,9 +678,31 @@ class DFlashModel(Qwen3Model):
+     def set_gguf_parameters(self):
+         super().set_gguf_parameters()
+ 
+-        block_size = self.hparams.get("block_size", 16)
+-        self.gguf_writer.add_block_size(block_size)
+         dflash_config = self.hparams.get("dflash_config", {})
++        block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16))
++        self.gguf_writer.add_block_size(block_size)
++
++        if "conv_kernel_size" in dflash_config:
++            self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"]))
++            self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"]))
++            self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"]))
++            self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"]))
++
++        output_multiplier = dflash_config.get(
++            "output_multiplier", self.hparams.get("output_multiplier")
++        )
++        if output_multiplier is not None:
++            self.gguf_writer.add_logit_scale(float(output_multiplier))
++        softcap = dflash_config.get(
++            "final_logit_softcapping", self.hparams.get("final_logit_softcapping")
++        )
++        if softcap is not None and float(softcap) > 0:
++            self.gguf_writer.add_final_logit_softcapping(float(softcap))
++        embedding_scale = dflash_config.get(
++            "input_embedding_scale", self.hparams.get("input_embedding_scale")
++        )
++        if embedding_scale is not None:
++            self.gguf_writer.add_embedding_scale(float(embedding_scale))
+ 
+         target_layer_ids = dflash_config.get("target_layer_ids", [])
+         if target_layer_ids:
+@@ -706,6 +728,12 @@ class DFlashModel(Qwen3Model):
+         if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True):
+             return
+ 
++        if name in (
++            "model.candidate_selector.predecessor_codebook",
++            "model.candidate_selector.successor_codebook",
++        ):
++            name += ".weight"
++
+         yield from super().modify_tensors(data_torch, name, bid)
+ 
+ 
+diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp
+index 487ae03ab..d074f270f 100644
+--- a/examples/speculative-simple/speculative-simple.cpp
++++ b/examples/speculative-simple/speculative-simple.cpp
+@@ -66,7 +66,8 @@ int main(int argc, char ** argv) {
+     llama_context * ctx_dft = params.speculative.draft.ctx_dft;
+ 
+     // check if the context supports partial sequence removal
+-    const bool use_ckpt_tgt = common_context_can_seq_rm(ctx_tgt) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
++    const auto ctx_tgt_seq_rm_type = common_context_can_seq_rm(ctx_tgt);
++    const bool use_ckpt_tgt = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
+     const bool use_ckpt_dft = common_context_can_seq_rm(ctx_dft) == COMMON_CONTEXT_SEQ_RM_TYPE_FULL;
+ 
+     if (use_ckpt_tgt) {
+@@ -152,6 +153,7 @@ int main(int argc, char ** argv) {
+     llama_batch batch_tgt = llama_batch_init(llama_n_batch(ctx_tgt), 0, 1);
+ 
+     llama_tokens draft;
++    std::vector<common_speculative_token_dist> dists;
+ 
+     common_prompt_checkpoint ckpt;
+ 
+@@ -192,6 +194,9 @@ int main(int argc, char ** argv) {
+                 /* .id_last    = */ id_last,
+                 /* .prompt     = */ &prompt_tgt,
+                 /* .result     = */ &draft, // output
++                /* .dists      = */ &dists,
++                /* .temperature = */ params.sampling.temp,
++                /* .seed       = */ common_sampler_get_seed(smpl.get()),
+             };
+             common_speculative_draft(spec);
+ 
+@@ -255,7 +260,12 @@ int main(int argc, char ** argv) {
+         // available logits from the batch and sample the next token until we run out of logits or the sampler
+         // disagrees with the draft
+         //
+-        auto ids = common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft);
++        const bool can_rollback =
++            ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART ||
++            (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && draft.size() <= llama_n_rs_seq(ctx_tgt));
++        auto ids = can_rollback && params.sampling.temp > 0.0f && dists.size() == draft.size()
++            ? common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft, dists)
++            : common_sampler_sample_and_accept_n(smpl.get(), ctx_tgt, draft);
+ 
+         //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str());
+ 
+@@ -268,6 +278,7 @@ int main(int argc, char ** argv) {
+             LOG_DBG("partial acceptance: %zu < %zu, restoring checkpoint\n", ids.size() - 1, n_draft);
+ 
+             draft = std::move(ids);
++            dists.clear();
+ 
+             {
+                 ckpt.load_tgt(ctx_tgt, seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY);
+@@ -325,6 +336,7 @@ int main(int argc, char ** argv) {
+ 
+         // clear the draft since it has been consumed
+         draft.clear();
++        dists.clear();
+ 
+         {
+             LOG_DBG("clear kv cache from any extra tokens, n_past = %d\n", n_past);
+diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
+index fad8d1fd8..3c81675c5 100644
+--- a/gguf-py/gguf/constants.py
++++ b/gguf-py/gguf/constants.py
+@@ -162,6 +162,10 @@ class Keys:
+         TARGET_LAYERS                     = "{arch}.target_layers"
+         TARGET_HIDDEN_SIZE                = "{arch}.target_hidden_size"
+         BLOCK_SIZE                        = "{arch}.block_size"
++        CONV_KERNEL_SIZE                  = "{arch}.conv_kernel_size"
++        CONV_GROUP_SIZE                   = "{arch}.conv_group_size"
++        SELECTOR_RANK                     = "{arch}.selector_rank"
++        SELECTOR_TOP_K                    = "{arch}.selector_top_k"
+         SAMPLE_FROM_ANCHOR                = "{arch}.sample_from_anchor"
+         NORM_BEFORE_RESIDUAL              = "{arch}.norm_before_residual"
+         NORM_BEFORE_FC                    = "{arch}.norm_before_fc"
+@@ -1135,6 +1139,13 @@ class MODEL_TENSOR(IntEnum):
+     DSPARK_MARKOV_W1       = auto()  # markov head: prev-token embed
+     DSPARK_MARKOV_W2       = auto()  # markov head: bias projection
+     DSPARK_CONF_PROJ       = auto()  # confidence head
++    DFLASH_ATTN_CONV_BASE  = auto()
++    DFLASH_ATTN_CONV_PROJ  = auto()
++    DFLASH_FFN_CONV_BASE   = auto()
++    DFLASH_FFN_CONV_PROJ   = auto()
++    DFLASH_SELECTOR_PREV   = auto()
++    DFLASH_SELECTOR_NEXT   = auto()
++    DFLASH_SELECTOR_HIDDEN = auto()
+     # lfm2 audio
+     A_ENC_NORM_CONV        = auto()
+     A_ENC_LINEAR_POS       = auto()
+@@ -1878,6 +1889,13 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
+     MODEL_TENSOR.DSPARK_MARKOV_W1:          "markov_w1",
+     MODEL_TENSOR.DSPARK_MARKOV_W2:          "markov_w2",
+     MODEL_TENSOR.DSPARK_CONF_PROJ:          "conf_proj",
++    MODEL_TENSOR.DFLASH_ATTN_CONV_BASE:      "blk.{bid}.attn_conv_base",
++    MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ:      "blk.{bid}.attn_conv_proj",
++    MODEL_TENSOR.DFLASH_FFN_CONV_BASE:       "blk.{bid}.ffn_conv_base",
++    MODEL_TENSOR.DFLASH_FFN_CONV_PROJ:       "blk.{bid}.ffn_conv_proj",
++    MODEL_TENSOR.DFLASH_SELECTOR_PREV:       "selector_predecessor",
++    MODEL_TENSOR.DFLASH_SELECTOR_NEXT:       "selector_successor",
++    MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN:     "selector_hidden",
+     MODEL_TENSOR.D2T:                       "d2t",
+ }
+ 
+@@ -4893,6 +4911,13 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
+         MODEL_TENSOR.DSPARK_MARKOV_W1,
+         MODEL_TENSOR.DSPARK_MARKOV_W2,
+         MODEL_TENSOR.DSPARK_CONF_PROJ,
++        MODEL_TENSOR.DFLASH_ATTN_CONV_BASE,
++        MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ,
++        MODEL_TENSOR.DFLASH_FFN_CONV_BASE,
++        MODEL_TENSOR.DFLASH_FFN_CONV_PROJ,
++        MODEL_TENSOR.DFLASH_SELECTOR_PREV,
++        MODEL_TENSOR.DFLASH_SELECTOR_NEXT,
++        MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN,
+     ],
+     MODEL_ARCH.MISTRAL4: [
+         MODEL_TENSOR.TOKEN_EMBD,
+diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py
+index 16ae9f999..b353fbf06 100644
+--- a/gguf-py/gguf/gguf_writer.py
++++ b/gguf-py/gguf/gguf_writer.py
+@@ -984,6 +984,18 @@ class GGUFWriter:
+     def add_block_size(self, value: int) -> None:
+         self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value)
+ 
++    def add_conv_kernel_size(self, value: int) -> None:
++        self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value)
++
++    def add_conv_group_size(self, value: int) -> None:
++        self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value)
++
++    def add_selector_rank(self, value: int) -> None:
++        self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value)
++
++    def add_selector_top_k(self, value: int) -> None:
++        self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value)
++
+     def add_sample_from_anchor(self, value: bool) -> None:
+         self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value)
+ 
+diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
+index a0571ccd3..6c4b3656d 100644
+--- a/gguf-py/gguf/tensor_mapping.py
++++ b/gguf-py/gguf/tensor_mapping.py
+@@ -1354,6 +1354,34 @@ class TensorNameMap:
+             "model.confidence_head.proj", # dspark
+         ),
+ 
++        MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: (
++            "model.layers.{bid}.attention_conv.base_kernel",
++        ),
++
++        MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: (
++            "model.layers.{bid}.attention_conv.kernel_projection",
++        ),
++
++        MODEL_TENSOR.DFLASH_FFN_CONV_BASE: (
++            "model.layers.{bid}.mlp_conv.base_kernel",
++        ),
++
++        MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: (
++            "model.layers.{bid}.mlp_conv.kernel_projection",
++        ),
++
++        MODEL_TENSOR.DFLASH_SELECTOR_PREV: (
++            "model.candidate_selector.predecessor_codebook",
++        ),
++
++        MODEL_TENSOR.DFLASH_SELECTOR_NEXT: (
++            "model.candidate_selector.successor_codebook",
++        ),
++
++        MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: (
++            "model.candidate_selector.hidden_projection",
++        ),
++
+         MODEL_TENSOR.CLS: (
+             "classifier",       # jina
+             "classifier.dense", # roberta
+diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
+index 955c2d796..e470d0b85 100644
+--- a/src/llama-arch.cpp
++++ b/src/llama-arch.cpp
+@@ -337,6 +337,11 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
+ 
+     { LLM_KV_TARGET_LAYERS,         "%s.target_layers"        },
+     { LLM_KV_TARGET_HIDDEN_SIZE,    "%s.target_hidden_size"   },
++    { LLM_KV_DFLASH_BLOCK_SIZE,       "%s.block_size"           },
++    { LLM_KV_DFLASH_CONV_KERNEL_SIZE, "%s.conv_kernel_size"    },
++    { LLM_KV_DFLASH_CONV_GROUP_SIZE,  "%s.conv_group_size"     },
++    { LLM_KV_DFLASH_SELECTOR_RANK,    "%s.selector_rank"       },
++    { LLM_KV_DFLASH_SELECTOR_TOP_K,   "%s.selector_top_k"      },
+     { LLM_KV_NORM_BEFORE_RESIDUAL,  "%s.norm_before_residual" },
+     { LLM_KV_NORM_BEFORE_FC,        "%s.norm_before_fc"       },
+ 
+@@ -647,6 +652,13 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
+     { LLM_TENSOR_DSPARK_MARKOV_W1,                       "markov_w1" },
+     { LLM_TENSOR_DSPARK_MARKOV_W2,                       "markov_w2" },
+     { LLM_TENSOR_DSPARK_CONF_PROJ,                       "conf_proj" },
++    { LLM_TENSOR_DFLASH_ATTN_CONV_BASE,                  "blk.%d.attn_conv_base" },
++    { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,                  "blk.%d.attn_conv_proj" },
++    { LLM_TENSOR_DFLASH_FFN_CONV_BASE,                   "blk.%d.ffn_conv_base" },
++    { LLM_TENSOR_DFLASH_FFN_CONV_PROJ,                   "blk.%d.ffn_conv_proj" },
++    { LLM_TENSOR_DFLASH_SELECTOR_PREV,                   "selector_predecessor" },
++    { LLM_TENSOR_DFLASH_SELECTOR_NEXT,                   "selector_successor" },
++    { LLM_TENSOR_DFLASH_SELECTOR_HIDDEN,                 "selector_hidden" },
+ };
+ 
+ // declare information about the model weight tensors:
+@@ -912,6 +924,13 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
+     {LLM_TENSOR_DSPARK_MARKOV_W1,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
+     {LLM_TENSOR_DSPARK_MARKOV_W2,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
+     {LLM_TENSOR_DSPARK_CONF_PROJ,           {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
++    {LLM_TENSOR_DFLASH_ATTN_CONV_BASE,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
++    {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,      {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
++    {LLM_TENSOR_DFLASH_FFN_CONV_BASE,       {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
++    {LLM_TENSOR_DFLASH_FFN_CONV_PROJ,       {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
++    {LLM_TENSOR_DFLASH_SELECTOR_PREV,       {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
++    {LLM_TENSOR_DFLASH_SELECTOR_NEXT,       {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_GET_ROWS}},
++    {LLM_TENSOR_DFLASH_SELECTOR_HIDDEN,     {LLM_TENSOR_LAYER_OUTPUT,    GGML_OP_MUL_MAT}},
+ };
+ 
+ LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {}
+diff --git a/src/llama-arch.h b/src/llama-arch.h
+index 48fe051a9..c0f218879 100644
+--- a/src/llama-arch.h
++++ b/src/llama-arch.h
+@@ -383,6 +383,11 @@ enum llm_kv {
+ 
+     LLM_KV_TARGET_LAYERS,
+     LLM_KV_TARGET_HIDDEN_SIZE,
++    LLM_KV_DFLASH_BLOCK_SIZE,
++    LLM_KV_DFLASH_CONV_KERNEL_SIZE,
++    LLM_KV_DFLASH_CONV_GROUP_SIZE,
++    LLM_KV_DFLASH_SELECTOR_RANK,
++    LLM_KV_DFLASH_SELECTOR_TOP_K,
+     LLM_KV_NORM_BEFORE_RESIDUAL,
+     LLM_KV_NORM_BEFORE_FC,
+ 
+@@ -655,6 +660,13 @@ enum llm_tensor {
+     LLM_TENSOR_DSPARK_MARKOV_W1,
+     LLM_TENSOR_DSPARK_MARKOV_W2,
+     LLM_TENSOR_DSPARK_CONF_PROJ,
++    LLM_TENSOR_DFLASH_ATTN_CONV_BASE,
++    LLM_TENSOR_DFLASH_ATTN_CONV_PROJ,
++    LLM_TENSOR_DFLASH_FFN_CONV_BASE,
++    LLM_TENSOR_DFLASH_FFN_CONV_PROJ,
++    LLM_TENSOR_DFLASH_SELECTOR_PREV,
++    LLM_TENSOR_DFLASH_SELECTOR_NEXT,
++    LLM_TENSOR_DFLASH_SELECTOR_HIDDEN,
+ };
+ 
+ 
+diff --git a/src/llama-context.cpp b/src/llama-context.cpp
+index 52f8d5367..51148b2ad 100644
+--- a/src/llama-context.cpp
++++ b/src/llama-context.cpp
+@@ -2314,6 +2314,12 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const {
+         }
+     }
+ 
++    if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) {
++        const uint32_t selector_tokens = std::min<uint32_t>(
++                n_tokens, model.hparams.dflash_block_size * cparams.n_seq_max);
++        res += 32*selector_tokens;
++    }
++
+     uint32_t n_sampling_nodes = 0;
+     uint32_t n_sampling_nodes_max = 0;
+     for (const auto & [seq_id, sampler] : sampling.samplers) {
+diff --git a/src/llama-graph.h b/src/llama-graph.h
+index 94324c745..605733f32 100644
+--- a/src/llama-graph.h
++++ b/src/llama-graph.h
+@@ -1326,6 +1326,8 @@ struct llm_graph_context {
+ 
+     void build_sampling() const;
+ 
++    virtual void build_post_sampling() const {}
++
+     //
+     // dense (out)
+     //
+diff --git a/src/llama-hparams.h b/src/llama-hparams.h
+index f6af36436..368ad7c70 100644
+--- a/src/llama-hparams.h
++++ b/src/llama-hparams.h
+@@ -218,6 +218,12 @@ struct llama_hparams {
+     // output embedding dimension (0 = use n_embd)
+     uint32_t n_embd_out_impl = 0;
+ 
++    uint32_t dflash_block_size       = 0;
++    uint32_t dflash_conv_kernel_size = 0;
++    uint32_t dflash_conv_group_size  = 0;
++    uint32_t dflash_selector_rank    = 0;
++    uint32_t dflash_selector_top_k   = 0;
++
+     // llama4 smallthinker
+     uint32_t n_moe_layer_step        = 0;
+     uint32_t n_no_rope_layer_step    = 4;
+diff --git a/src/llama-model.cpp b/src/llama-model.cpp
+index 3759c8625..5dbd22f63 100644
+--- a/src/llama-model.cpp
++++ b/src/llama-model.cpp
+@@ -2462,6 +2462,7 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const {
+ 
+     // add backend sampling layers (if any)
+     llm->build_sampling();
++    llm->build_post_sampling();
+ 
+     // if the gguf model was converted with --sentence-transformers-dense-modules
+     // there will be two additional dense projection layers
+diff --git a/src/llama-model.h b/src/llama-model.h
+index 4412ef08e..e2cd04946 100644
+--- a/src/llama-model.h
++++ b/src/llama-model.h
+@@ -362,6 +362,11 @@ struct llama_layer {
+     struct ggml_tensor * ffn_exp_probs_b = nullptr;
+     struct ggml_tensor * ffn_gate_tid2eid = nullptr;
+ 
++    struct ggml_tensor * dflash_attn_conv_base = nullptr;
++    struct ggml_tensor * dflash_attn_conv_proj = nullptr;
++    struct ggml_tensor * dflash_ffn_conv_base  = nullptr;
++    struct ggml_tensor * dflash_ffn_conv_proj  = nullptr;
++
+     // mamba proj
+     struct ggml_tensor * ssm_in  = nullptr;
+     struct ggml_tensor * ssm_x   = nullptr;
+@@ -648,6 +653,10 @@ struct llama_model {
+     struct ggml_tensor * dspark_conf_proj   = nullptr;
+     struct ggml_tensor * dspark_conf_proj_b = nullptr;
+ 
++    struct ggml_tensor * dflash_selector_prev   = nullptr;
++    struct ggml_tensor * dflash_selector_next   = nullptr;
++    struct ggml_tensor * dflash_selector_hidden = nullptr;
++
+     // unified vector to store target-model extracted layer ids in eagle3, dflash, etc.
+     std::vector<int32_t> target_layer_ids;
+ 
+diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
+index 5b70a5179..776d34dd2 100644
+--- a/src/models/dflash.cpp
++++ b/src/models/dflash.cpp
+@@ -7,6 +7,16 @@
+ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
+ 
+     ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
++    ml.get_key(LLM_KV_LOGIT_SCALE,                  hparams.f_logit_scale, false);
++    hparams.f_final_logit_softcapping = 0.0f;
++    ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING,      hparams.f_final_logit_softcapping, false);
++    ml.get_key(LLM_KV_EMBEDDING_SCALE,              hparams.f_embedding_scale, false);
++
++    ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE,        hparams.dflash_block_size,       false);
++    ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false);
++    ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE,  hparams.dflash_conv_group_size,  false);
++    ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK,    hparams.dflash_selector_rank,    false);
++    ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K,   hparams.dflash_selector_top_k,   false);
+ 
+     if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) {
+         throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata");
+@@ -112,6 +122,29 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
+         LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
+     }
+ 
++    const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight");
++    if (selector_meta) {
++        const int64_t rank = hparams.dflash_selector_rank;
++        if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 ||
++                hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) {
++            throw std::runtime_error("DFlash2 model is missing conv/selector metadata");
++        }
++        if (n_embd % hparams.dflash_conv_group_size != 0) {
++            throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size");
++        }
++        if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) {
++            throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice");
++        }
++
++        dflash_selector_prev   = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV,   "weight"), { rank, n_vocab }, 0);
++        dflash_selector_next   = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT,   "weight"), { rank, n_vocab }, 0);
++        dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0);
++
++        LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__,
++                hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size,
++                hparams.dflash_selector_rank, hparams.dflash_selector_top_k);
++    }
++
+     fc              = create_tensor(tn(LLM_TENSOR_FC,              "weight"), { n_embd_inp, n_embd }, 0);
+     fc_s            = create_tensor(tn(LLM_TENSOR_FC,              "scale"),  { 1 }, TENSOR_NOT_REQUIRED);
+     output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
+@@ -187,6 +220,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
+         layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
+         layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
+         layer.ffn_up   = create_tensor(tn(LLM_TENSOR_FFN_UP,   "weight", i), { n_embd, n_ff }, 0);
++
++        if (selector_meta) {
++            const int64_t kernel = hparams.dflash_conv_kernel_size;
++            const int64_t groups = n_embd / hparams.dflash_conv_group_size;
++            const int64_t projected = 2 * kernel * groups;
++            layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
++            layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0);
++            layer.dflash_ffn_conv_base  = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0);
++            layer.dflash_ffn_conv_proj  = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ,  "weight", i), { n_embd, projected }, 0);
++        }
+     }
+ }
+ 
+@@ -222,7 +265,7 @@ ggml_tensor * llama_model_dflash::graph<true>::build_inp_embd_enc() const {
+ 
+ // DFlash Encoder: processes target model features through feature fusion layer
+ template <>
+-llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
++llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model) {
+     ggml_tensor * cur = build_inp_embd_enc();
+ 
+     cur = build_lora_mm(model.fc, cur, model.fc_s);
+@@ -345,11 +388,69 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
+     ggml_build_forward_expand(g.gf, out);
+ }
+ 
++static ggml_tensor * build_dflash2_conv(
++        llm_graph_context & g,
++        ggml_tensor * hidden,
++        ggml_tensor * dynamic,
++        ggml_tensor * base,
++        int side) {
++    const auto & hparams = g.hparams;
++    const int64_t hidden_size = hidden->ne[0];
++    const int64_t n_tokens    = hidden->ne[1];
++    const int64_t n_blocks    = g.ubatch.n_seqs_unq;
++    const int64_t kernel_size = hparams.dflash_conv_kernel_size;
++    const int64_t group_size  = hparams.dflash_conv_group_size;
++    const int64_t n_groups    = hidden_size / group_size;
++
++    GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
++    GGML_ASSERT(dynamic && base && side >= 0 && side < 2);
++
++    const int64_t block_size = n_tokens / n_blocks;
++    ggml_context * ctx0 = g.ctx0;
++    hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens);
++    dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens);
++    ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks);
++    ggml_tensor * grouped = ggml_reshape_3d(ctx0, hidden, group_size, n_groups, n_tokens);
++    ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens);
++    ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens,
++            coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]);
++
++    ggml_tensor * result = nullptr;
++    for (int64_t tap = 0; tap < kernel_size; ++tap) {
++        ggml_tensor * values = blocks;
++        if (tap > 0) {
++            ggml_tensor * zeros = ggml_fill(ctx0,
++                    ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f);
++            if (tap < block_size) {
++                ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks,
++                        blocks->nb[1], blocks->nb[2], 0);
++                values = ggml_concat(ctx0, zeros, previous, 1);
++            } else {
++                values = zeros;
++            }
++        }
++        values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens);
++
++        ggml_tensor * coeff = ggml_view_2d(ctx0, coeffs_side, n_groups, n_tokens,
++                coeffs_side->nb[2], tap * coeffs_side->nb[1]);
++        coeff = ggml_cont(ctx0, coeff);
++        coeff = ggml_reshape_3d(ctx0, coeff, 1, n_groups, n_tokens);
++        coeff = ggml_reshape_2d(ctx0, ggml_repeat(ctx0, coeff, grouped), hidden_size, n_tokens);
++
++        ggml_tensor * base_tap = ggml_view_1d(ctx0, base, hidden_size,
++                tap * base->nb[1] + side * base->nb[2]);
++        ggml_tensor * weight = ggml_add(ctx0, coeff, ggml_repeat(ctx0, base_tap, hidden));
++        ggml_tensor * term = ggml_mul(ctx0, weight, values);
++        result = result ? ggml_add(ctx0, result, term) : term;
++    }
++    return result;
++}
++
+ // DFlash decoder, dual-mode by batch type:
+ //   * embd batch  -> fused target features: project + inject K/V into the cache.
+ //   * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens
+ template <>
+-llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
++llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params), model(model) {
+     const int64_t n_embd_head = hparams.n_embd_head_v();
+ 
+     GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
+@@ -449,10 +550,14 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
+ 
+     inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+     ggml_set_input(inp->tokens);
++    res->t_inp_tokens = inp->tokens;
+ 
+     ggml_tensor * inp_tokens = inp->tokens;
+ 
+     ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
++    if (hparams.f_embedding_scale != 0.0f) {
++        inpL = ggml_scale(ctx0, inpL, hparams.f_embedding_scale);
++    }
+     cb(inpL, "inp_noise_embd", -1);
+ 
+     res->add_input(std::move(inp));
+@@ -463,6 +568,13 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
+         ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il);
+         cb(noise_norm, "noise_norm", il);
+ 
++        ggml_tensor * attn_dynamic = nullptr;
++        if (layer.dflash_attn_conv_proj) {
++            attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm);
++            noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0);
++            cb(noise_norm, "attn_conv_in", il);
++        }
++
+         ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm);
+         ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm);
+         ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm);
+@@ -493,12 +605,24 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
+             ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il)
+             : build_attn(inp_attn,      layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il);
+ 
++        if (attn_dynamic) {
++            cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1);
++            cb(cur, "attn_conv_out", il);
++        }
++
+         ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
+         cb(ffn_inp, "ffn_inp", il);
+ 
+         cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+         cb(cur, "ffn_norm", il);
+ 
++        ggml_tensor * ffn_dynamic = nullptr;
++        if (layer.dflash_ffn_conv_proj) {
++            ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur);
++            cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0);
++            cb(cur, "ffn_conv_in", il);
++        }
++
+         cur = build_ffn(cur,
+                 layer.ffn_up,   NULL, layer.ffn_up_s,
+                 layer.ffn_gate, NULL, layer.ffn_gate_s,
+@@ -507,6 +631,11 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
+                 LLM_FFN_SILU, LLM_FFN_PAR, il);
+         cb(cur, "ffn_out", il);
+ 
++        if (ffn_dynamic) {
++            cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1);
++            cb(cur, "ffn_conv_out", il);
++        }
++
+         cur = ggml_add(ctx0, cur, ffn_inp);
+         cb(cur, "l_out", il);
+ 
+@@ -531,6 +660,15 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
+ 
+     cur = build_lora_mm(output, cur, output_s);
+ 
++    if (hparams.f_logit_scale != 0.0f) {
++        cur = ggml_scale(ctx0, cur, hparams.f_logit_scale);
++    }
++    if (hparams.f_final_logit_softcapping > 0.0f) {
++        cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
++        cur = ggml_tanh(ctx0, cur);
++        cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
++    }
++
+     // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t
+     if (model.d2t) {
+         const int64_t n_draft_vocab = cur->ne[0];
+@@ -557,6 +695,98 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
+     }
+ }
+ 
++template <bool is_enc>
++void llama_model_dflash::graph<is_enc>::build_post_sampling() const {
++    if constexpr (is_enc) {
++        return;
++    }
++
++    if (!model.dflash_selector_hidden || !res->t_logits) {
++        return;
++    }
++
++    const int64_t top_k    = hparams.dflash_selector_top_k;
++    const int64_t rank     = hparams.dflash_selector_rank;
++    const int64_t n_blocks = ubatch.n_seqs_unq;
++    GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0);
++    GGML_ASSERT(res->t_logits->ne[1] == n_tokens);
++    ggml_tensor * tokens = res->get_inp_tokens();
++    if (!tokens) {
++        return;
++    }
++
++    const int64_t tokens_per_block = n_tokens / n_blocks;
++    const int64_t block_size = std::min<int64_t>(tokens_per_block, hparams.dflash_block_size);
++    ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k);
++    ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens);
++    ggml_tensor * unary = ggml_reshape_2d(ctx0,
++            ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens);
++
++    std::vector<ggml_tensor *> candidate_ids(block_size);
++    std::vector<ggml_tensor *> unary_logits(block_size);
++    for (int64_t pos = 1; pos < block_size; ++pos) {
++        candidate_ids[pos] = ggml_cont_2d(ctx0,
++                ggml_view_2d(ctx0, candidates, top_k, n_blocks,
++                    tokens_per_block * candidates->nb[1], pos * candidates->nb[1]),
++                top_k, n_blocks);
++        unary_logits[pos] = ggml_cont_2d(ctx0,
++                ggml_view_2d(ctx0, unary, top_k, n_blocks,
++                    tokens_per_block * unary->nb[1], pos * unary->nb[1]),
++                top_k, n_blocks);
++    }
++
++    ggml_tensor * hidden = build_lora_mm(model.dflash_selector_hidden, res->t_embd);
++
++    ggml_tensor * anchor_ids = ggml_view_2d(ctx0, tokens, 1, n_blocks,
++            tokens_per_block * tokens->nb[0], 0);
++    anchor_ids = ggml_cont_1d(ctx0, anchor_ids, n_blocks);
++
++    ggml_tensor * packed = ggml_fill(ctx0,
++            ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f);
++
++    for (int64_t pos = 1; pos < block_size; ++pos) {
++        ggml_tensor * ids = candidate_ids[pos];
++        ggml_tensor * unary = unary_logits[pos];
++        ggml_tensor * successor = ggml_get_rows(ctx0, model.dflash_selector_next,
++                ggml_reshape_1d(ctx0, ids, top_k * n_blocks));
++        successor = ggml_reshape_3d(ctx0, successor, rank, top_k, n_blocks);
++
++        ggml_tensor * hidden_pos = ggml_cont(ctx0, ggml_view_2d(ctx0, hidden, rank, n_blocks,
++                tokens_per_block * hidden->nb[1], pos * hidden->nb[1]));
++        hidden_pos = ggml_reshape_3d(ctx0, hidden_pos, rank, 1, n_blocks);
++
++        ggml_tensor * predecessor;
++        if (pos == 1) {
++            predecessor = ggml_get_rows(ctx0, model.dflash_selector_prev, anchor_ids);
++            predecessor = ggml_reshape_3d(ctx0, predecessor, rank, 1, n_blocks);
++        } else {
++            predecessor = ggml_get_rows(ctx0, model.dflash_selector_prev,
++                    ggml_reshape_1d(ctx0, candidate_ids[pos - 1], top_k * n_blocks));
++            predecessor = ggml_reshape_3d(ctx0, predecessor, rank, top_k, n_blocks);
++        }
++
++        ggml_tensor * conditioned = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, hidden_pos, predecessor));
++        ggml_tensor * scores = ggml_mul_mat(ctx0, successor, conditioned);
++        if (pos == 1) {
++            scores = ggml_repeat_4d(ctx0, scores, top_k, top_k, n_blocks, 1);
++        }
++        ggml_tensor * unary_3d = ggml_reshape_3d(ctx0, unary, top_k, 1, n_blocks);
++        scores = ggml_add(ctx0, scores, ggml_repeat(ctx0, unary_3d, scores));
++
++        ggml_tensor * row = ggml_concat(ctx0,
++                ggml_cast(ctx0, ids, GGML_TYPE_F32),
++                ggml_reshape_2d(ctx0, scores, top_k * top_k, n_blocks), 0);
++        row = ggml_pad(ctx0, row, n_embd - row->ne[0], 0, 0, 0);
++        row = ggml_reshape_3d(ctx0, row, n_embd, 1, n_blocks);
++        packed = ggml_concat(ctx0, packed, row, 1);
++    }
++
++    packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks);
++    cb(packed, "dflash2_lattice", -1);
++    res->t_h_nextn = packed;
++    ggml_build_forward_expand(gf, packed);
++}
++
+ // DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above):
+ //   * embd batch  -> project main_x through each stage's wkv and inject K into the ring cache
+ //   * token batch -> noise block through 3 full DSV4 stages (hc + MLA + MoE), markov + confidence heads
+diff --git a/src/models/models.h b/src/models/models.h
+index 1dd30dfd1..f4a72fcba 100644
+--- a/src/models/models.h
++++ b/src/models/models.h
+@@ -1342,9 +1342,12 @@ struct llama_model_dflash : public llama_model_base {
+ 
+     template <bool is_enc>
+     struct graph : public llm_graph_context {
++        const llama_model & model;
++
+         graph(const llama_model & model, const llm_graph_params & params);
+ 
+         ggml_tensor * build_inp_embd_enc() const;
++        void build_post_sampling() const override;
+     };
+ 
+     struct graph_dsv4 : public llama_model_deepseek4::graph {
+diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
+index 21ff78394..84665ae86 100644
+--- a/tools/server/server-context.cpp
++++ b/tools/server/server-context.cpp
+@@ -209,6 +209,7 @@ struct server_slot {
+     common_speculative * spec;
+ 
+     llama_tokens spec_draft;
++    std::vector<common_speculative_token_dist> spec_dists;
+     llama_tokens spec_prompt;
+     std::vector<int32_t> spec_i_batch;
+     common_prompt_checkpoint spec_ckpt;
+@@ -337,6 +338,7 @@ struct server_slot {
+ 
+         if (can_speculate()) {
+             spec_draft.clear();
++            spec_dists.clear();
+             spec_i_batch.clear();
+             spec_ckpt.clear();
+         }
+@@ -2961,6 +2963,9 @@ private:
+                             /* .id_last  = */ slot.sampled,
+                             /* .prompt   = */ &slot.spec_prompt,
+                             /* .result   = */ &slot.spec_draft,
++                            /* .dists    = */ &slot.spec_dists,
++                            /* .temperature = */ slot.task->params.sampling.temp,
++                            /* .seed     = */ common_sampler_get_seed(slot.smpl.get()),
+                         };
+ 
+                         drafting.push_back(&slot);
+@@ -3840,7 +3845,13 @@ private:
+                 common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get()));
+ 
+                 GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1);
+-                auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
++                const bool can_rollback =
++                    ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART ||
++                    (ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_RS && n_draft <= llama_n_rs_seq(ctx_tgt));
++                auto accepted = can_rollback && slot.task->params.sampling.temp > 0.0f &&
++                                slot.spec_dists.size() == slot.spec_draft.size()
++                    ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, slot.spec_dists)
++                    : common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft);
+                 slot.spec_i_batch.clear();
+ 
+                 GGML_ASSERT(accepted.size() >= 1);
+@@ -3861,6 +3872,7 @@ private:
+                         // partial acceptance is not supported by the context -> truncate the draft and restore the state
+                         slot.spec_is_replay = true;
+                         slot.spec_draft = std::move(accepted);
++                        slot.spec_dists.clear();
+ 
+                         const auto & ckpt = slot.spec_ckpt;
+ 
+@@ -3888,6 +3900,7 @@ private:
+                 common_speculative_accept(spec.get(), slot.id, accepted.size() - 1);
+ 
+                 slot.spec_draft = std::move(accepted);
++                slot.spec_dists.clear();
+             }
+ 
+             const auto ids = std::move(slot.spec_draft);