Lumpiasty/llama.cpp
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

spec: add DSpark speculative decoding (#25173)

Browse Source 
* spec: add DSpark speculative decoding

DSpark (DeepSpec, 2026) on top of the merged DFlash drafter. It reuses the
DFlash encoder/decoder graph, target feature extraction and KV-cache injection,
and the verify/accept path unchanged; the draft model is a new "dspark" arch
adding a low-rank Markov head (markov_w1/w2) and an optional (unused here)
confidence head. No new public APIs.

The proposal is the only change: the block is anchor-first (position 0 already
predicts the first draft) and the decoder graph applies a semi-autoregressive,
previous-token conditioned logit bias in-graph, chained per block position:

  logits'(i) = logits(i) + markov_w2 . markov_w1[prev(i)]
  prev(0)    = the block's anchor token, prev(i>0) = argmax(logits'(i-1))

vectorized across all blocks in the batch; the anchors are fed through a
dedicated graph input (token 0 of every block). Greedy stays lossless
(verify unchanged, same as DFlash).

- new arch "dspark" (llama_model_dspark : llama_model_dflash, reuses the graph,
  loads the markov/confidence tensors; shares the target's embed/lm_head).
- Qwen3DSparkModel converter.
- new spec type "draft-dspark" (common_speculative_impl_draft_dspark :
  common_speculative_impl_draft_dflash, overrides draft() only: submits whole
  anchor-first blocks and greedily reads back the biased logits).

* spec: read draft block size in the dflash impl

* docs: add DSpark section to speculative.md

* spec: keep dspark block size read in the dspark impl

* dspark : add TODOs for incomplete parts

- confidence head is loaded but not used yet
- confidence-scheduled prefix pruning is not implemented
- the in-graph Markov chain is greedy-only
- only Qwen3 backbones are supported for now (also noted in docs)

* spec: fold DSpark into the DFlash arch

Address review: drop LLM_ARCH_DSPARK and the dspark.block_size /
markov_rank GGUF keys. A DSpark draft now converts to a DFlash GGUF;
the Markov head tensors are detected by presence (like eagle3 d2t),
block_size is read from the existing dflash.block_size key, and the
block anchors are taken as a strided view of the decoder's token
input instead of a separate graph input.

* spec: add confidence-based draft pruning for DSpark

The DSpark confidence head predicts per-position acceptance of the
drafted block. --spec-draft-conf-min truncates the block at the first
position below the threshold (default 0 = disabled).

* fold the dspark impl into dflash, selected by spec type

* address review comments

* dspark: clean up and improve naming

* update readme

* remove trailing whitespace

* dflash: draft full n_max blocks, defer dp.n_max to the central truncation

The DSpark markov head views the draft batch as a uniform [n_seqs x block]
grid, but the per-seq dp.n_max clamp could produce blocks of different
sizes, silently corrupting the strided views and the resulting logits.

Drop the clamp and always draft the full n_max block for every sequence:
dp.n_max is already enforced by the central truncation in
common_speculative_draft(), the same way eagle3 handles it.

Co-authored-by: Zaire404 <3147879462@qq.com>

* dflash: assert the markov head block-uniformity invariant, require the conf head

With the draft batch always submitting equal-size n_max blocks, a
non-divisible token count can only mean the batch was split across
ubatches or a caller broke the layout - fail loudly instead of silently
dropping the markov bias. The block_drafts > block_size early return
stays: worst-case graph reserve passes legitimately build with
n_seq_tokens > block_size.

Also make conf_proj required when the markov head is present: the
confidence head is part of the DSpark checkpoint format, and a missing
head would otherwise leave --spec-draft-conf-min silently reading stale
embeddings instead of confidences.

Co-authored-by: Zaire404 <3147879462@qq.com>

* dspark: fold conf_min into p_min

p_min and conf_min express the same thing - the minimum predicted
survival probability for a drafted position - differing only in how the
estimate is obtained: token probability for regular drafters, the
trained confidence head for DSpark. The DSpark readback never used
p_min, so reuse it for the confidence threshold and drop the separate
--spec-draft-conf-min flag. Both defaulted to 0 (disabled), so behavior
is unchanged.

Co-authored-by: Zaire404 <3147879462@qq.com>

* dflash: note the confidence broadcast workaround

Requested in review: the ggml_repeat only adapts the [1, n_tok]
confidences to the n_embd-wide embd_nextn transport so that
llama_get_embeddings_nextn can be reused - not a placeholder.

Co-authored-by: Zaire404 <3147879462@qq.com>

* cont : clarify

[no ci]

---------

Co-authored-by: Ruixiang Wang <wangruixiang07@outlook.com>
Co-authored-by: Zaire404 <3147879462@qq.com>
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
王金旭
2026-07-28 14:43:27 +03:00
committed by GitHub
co-authored by Zaire404 Ruixiang Wang Georgi Gerganov
parent 6ba5ef2470
commit 84075273c8
14 changed files with 287 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
common/speculative.cpp
+76 -30
View File
@@ -34,6 +34,7 @@ const std::map<std::string, common_speculative_type> common_speculative_type_fro
{"draft-eagle3", COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3},
{"draft-mtp", COMMON_SPECULATIVE_TYPE_DRAFT_MTP},
{"draft-dflash", COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH},
{"draft-dspark", COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK},
{"ngram-simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE},
{"ngram-map-k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K},
{"ngram-map-k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V},
@@ -928,15 +929,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
int32_t block_size = 0;
llama_token mask_token_id = 0;
// draft-dspark: the draft carries a Markov head and uses an anchor-first block layout
const bool is_dspark;
const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices
uint32_t target_layer_ids_n = 0;
// scratch buffer for concatenated target features [n_tokens, n_embd_enc]
std::vector<float> features_buf;
common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq)
: common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, n_seq)
common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq,
common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)
: common_speculative_impl(type, n_seq)
, params(params.draft)
, is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)
{
auto * ctx_tgt = this->params.ctx_tgt;
auto * ctx_dft = this->params.ctx_dft;
@@ -963,16 +969,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
}
mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft));
LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__);
LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str());
LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min);
LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u\n", __func__, block_size, mask_token_id, target_layer_ids_n);
// DFlash input is [id_last, <mask> * (block_size-1)], so it can draft at most block_size-1 tokens per step
if (this->params.n_max > block_size - 1 || this->params.n_min > block_size - 1) {
LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained DFlash block size %d -- clamping to %d\n",
__func__, this->params.n_max, this->params.n_min, block_size, block_size - 1);
this->params.n_max = std::min(this->params.n_max, block_size - 1);
this->params.n_min = std::min(this->params.n_min, block_size - 1);
// DFlash input is [id_last, <mask> * (block_size-1)]: in-place denoising yields at most
// block_size-1 draft tokens, DSpark yield a full block_size draft tokens
const int32_t n_draft_max = is_dspark ? block_size : block_size - 1;
if (this->params.n_max > n_draft_max || this->params.n_min > n_draft_max) {
LOG_WRN("%s: requested draft size (n_max=%d, n_min=%d) exceeds the trained block size %d -- clamping to %d\n",
__func__, this->params.n_max, this->params.n_min, block_size, n_draft_max);
this->params.n_max = std::min(this->params.n_max, n_draft_max);
this->params.n_min = std::min(this->params.n_min, n_draft_max);
}
batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq);
@@ -1136,12 +1144,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
const int32_t n = (int32_t) dp.n_past;
int32_t n_draft = params.n_max;
if (dp.n_max > 0) {
n_draft = std::min(n_draft, dp.n_max);
}
const int32_t n_draft = params.n_max;
const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * <mask>
const int32_t n_block_tokens = n_draft + (is_dspark ? 0 : 1);
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) {
@@ -1173,27 +1178,57 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl {
auto & result = *dp.result;
// greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
for (int32_t i = 1; i < n_block_tokens; ++i) {
common_sampler_sample(smpl, ctx_dft, beg + i, true);
if (is_dspark) {
// DSpark predicts the next token from position 0 and optionally truncates
// at the first position below the confidence threshold.
const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr;
const auto * cur_p = common_sampler_get_candidates(smpl, true);
for (int32_t i = 0; i < n_block_tokens; ++i) {
const int32_t idx = beg + i;
for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) {
break;
}
common_sampler_sample(smpl, ctx_dft, idx, true);
const auto * cur_p = common_sampler_get_candidates(smpl, true);
for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
seq_id, k, i, cur_p->data[k].id, cur_p->data[k].p,
common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
}
const llama_token id = cur_p->data[0].id;
common_sampler_accept(smpl, id, true);
result.push_back(id);
}
} else {
// greedily read the predicted block at this sequence's noise positions 1..n_block_tokens-1
for (int32_t i = 1; i < n_block_tokens; ++i) {
common_sampler_sample(smpl, ctx_dft, beg + i, true);
const llama_token id = cur_p->data[0].id;
const auto * cur_p = common_sampler_get_candidates(smpl, true);
if (cur_p->data[0].p < params.p_min) {
break;
for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) {
LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n",
seq_id, k, i - 1, cur_p->data[k].id, cur_p->data[k].p,
common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str());
}
const llama_token id = cur_p->data[0].id;
if (cur_p->data[0].p < params.p_min) {
break;
}
common_sampler_accept(smpl, id, true);
result.push_back(id);
}
common_sampler_accept(smpl, id, true);
result.push_back(id);
}
if (result.size() < (size_t) params.n_min) {
@@ -2155,6 +2190,7 @@ std::string common_speculative_type_to_str(common_speculative_type type) {
case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3: return "draft-eagle3";
case COMMON_SPECULATIVE_TYPE_DRAFT_MTP: return "draft-mtp";
case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH: return "draft-dflash";
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: return "draft-dspark";
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram-simple";
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram-map-k";
case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram-map-k4v";
@@ -2208,6 +2244,7 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) {
case COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3:
case COMMON_SPECULATIVE_TYPE_DRAFT_MTP:
case COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH:
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK:
n_max = std::max(n_max, std::max(0, spec->draft.n_max));
break;
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
@@ -2352,6 +2389,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
bool has_draft_eagle3 = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3)) && params.draft.ctx_dft != nullptr;
bool has_draft_mtp = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_MTP)) && params.draft.ctx_dft != nullptr;
bool has_draft_dflash = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH)) && params.draft.ctx_dft != nullptr;
bool has_draft_dspark = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)) && params.draft.ctx_dft != nullptr;
@@ -2362,7 +2400,7 @@ common_speculative * common_speculative_init(common_params_speculative & params,
bool has_ngram_mod = (enabled_configs & (1u << COMMON_SPECULATIVE_TYPE_NGRAM_MOD));
// when adding a new type - update here the logic above
static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 10);
static_assert(COMMON_SPECULATIVE_TYPE_COUNT == 11);
// this list here defines the priority of the speculators
// the one with highest priority are listed first
@@ -2395,6 +2433,9 @@ common_speculative * common_speculative_init(common_params_speculative & params,
if (has_draft_dflash) {
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, params));
}
if (has_draft_dspark) {
configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, params));
}
}
std::vector<std::unique_ptr<common_speculative_impl>> impls = {};
@@ -2419,6 +2460,11 @@ common_speculative * common_speculative_init(common_params_speculative & params,
impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(config.params, n_seq));
break;
}
case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: {
impls.push_back(std::make_unique<common_speculative_impl_draft_dflash>(
config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK));
break;
}
case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: {
common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple);