spec: add DSpark speculative decoding (#25173)
* 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:
王金旭
co-authored by
Zaire404
Ruixiang Wang
Georgi Gerganov
Zaire404
Ruixiang Wang
Georgi Gerganov
parent
6ba5ef2470
commit
84075273c8
@@ -616,6 +616,9 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
|
||||
{ LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" },
|
||||
{ LLM_TENSOR_FC, "fc" },
|
||||
{ LLM_TENSOR_D2T, "d2t" },
|
||||
{ LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
|
||||
{ LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
|
||||
{ LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
|
||||
};
|
||||
|
||||
// declare information about the model weight tensors:
|
||||
@@ -870,6 +873,10 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
|
||||
// eagle3
|
||||
{LLM_TENSOR_FC, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
|
||||
{LLM_TENSOR_D2T, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
|
||||
// dspark
|
||||
{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_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {}
|
||||
|
||||
@@ -624,6 +624,9 @@ enum llm_tensor {
|
||||
LLM_TENSOR_MASKED_EMBD_ORDERING,
|
||||
LLM_TENSOR_FC,
|
||||
LLM_TENSOR_D2T,
|
||||
LLM_TENSOR_DSPARK_MARKOV_W1,
|
||||
LLM_TENSOR_DSPARK_MARKOV_W2,
|
||||
LLM_TENSOR_DSPARK_CONF_PROJ,
|
||||
};
|
||||
|
||||
|
||||
|
||||
@@ -606,6 +606,12 @@ struct llama_model {
|
||||
struct ggml_tensor * fc = nullptr; // feature fusion layer
|
||||
struct ggml_tensor * d2t = nullptr; // draft to target vocabulary mapping
|
||||
|
||||
// dspark
|
||||
struct ggml_tensor * dspark_markov_w1 = nullptr;
|
||||
struct ggml_tensor * dspark_markov_w2 = nullptr;
|
||||
struct ggml_tensor * dspark_conf_proj = nullptr;
|
||||
struct ggml_tensor * dspark_conf_proj_b = nullptr;
|
||||
|
||||
// unified vector to store target-model extracted layer ids in eagle3, dflash, etc.
|
||||
std::vector<int32_t> target_layer_ids;
|
||||
|
||||
|
||||
@@ -37,6 +37,23 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
|
||||
|
||||
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
|
||||
|
||||
// DSpark = DFlash + a semi-autoregressive Markov head and Confidence head
|
||||
//
|
||||
// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4)
|
||||
// need their own conversion path and graph tweaks
|
||||
const struct ggml_tensor * markov_meta = ml->get_tensor_meta("markov_w1.weight");
|
||||
if (markov_meta) {
|
||||
const int64_t dspark_markov_rank = markov_meta->ne[0];
|
||||
|
||||
dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0);
|
||||
dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0);
|
||||
|
||||
dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0);
|
||||
dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);
|
||||
|
||||
LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
|
||||
}
|
||||
|
||||
fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0);
|
||||
output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc)
|
||||
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); // decoder final norm
|
||||
@@ -105,6 +122,94 @@ llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_grap
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
}
|
||||
|
||||
// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position
|
||||
static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
|
||||
ggml_context * ctx0 = g.ctx0;
|
||||
auto & res = g.res;
|
||||
|
||||
ggml_tensor * w1 = model.dspark_markov_w1;
|
||||
ggml_tensor * w2 = model.dspark_markov_w2;
|
||||
GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded");
|
||||
|
||||
ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]
|
||||
const int64_t n_vocab = base->ne[0];
|
||||
const int64_t n_tok = base->ne[1];
|
||||
|
||||
const auto it = model.gguf_kv.find("dflash.block_size");
|
||||
GGML_ASSERT(it != model.gguf_kv.end() && "DSpark draft requires 'dflash.block_size' in GGUF metadata");
|
||||
const int64_t block_size = std::stoi(it->second);
|
||||
GGML_ASSERT(block_size > 0);
|
||||
|
||||
const int64_t n_blocks = g.ubatch.n_seqs_unq;
|
||||
GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks");
|
||||
// runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size
|
||||
const int64_t block_drafts = n_tok / n_blocks;
|
||||
if (block_drafts > block_size) {
|
||||
return;
|
||||
}
|
||||
|
||||
// anchor (committed last) token of every block: token 0 of each block, i.e. a strided view
|
||||
const size_t token_stride = (size_t) block_drafts * tokens->nb[0];
|
||||
const size_t base_stride = (size_t) block_drafts * base->nb[1];
|
||||
|
||||
ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0);
|
||||
prev = ggml_cont_1d(ctx0, prev, n_blocks);
|
||||
|
||||
// confidence head input: predicts per-position acceptance
|
||||
ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok]
|
||||
|
||||
ggml_tensor * cat = nullptr;
|
||||
ggml_tensor * cat_conf = nullptr;
|
||||
|
||||
// TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
|
||||
// token pick, not the Markov conditioning path
|
||||
for (int64_t i = 0; i < block_drafts; ++i) {
|
||||
ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
|
||||
ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab, n_blocks]
|
||||
|
||||
// position i of every block: strided view [n_vocab, n_blocks]
|
||||
ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]);
|
||||
ggml_tensor * col = ggml_add(ctx0, base_i, bias);
|
||||
|
||||
cat = cat ? ggml_concat(ctx0, cat, col, 1) : col;
|
||||
|
||||
// conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks]
|
||||
ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
|
||||
(size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
|
||||
ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);
|
||||
ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
|
||||
if (model.dspark_conf_proj_b) {
|
||||
conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
|
||||
}
|
||||
conf = ggml_sigmoid(ctx0, conf);
|
||||
|
||||
cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf;
|
||||
|
||||
if (i + 1 < block_drafts) {
|
||||
prev = ggml_argmax(ctx0, col);
|
||||
}
|
||||
}
|
||||
|
||||
// cat is position-major; restore ubatch block-major order
|
||||
ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, block_drafts);
|
||||
out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks]
|
||||
out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok);
|
||||
|
||||
{
|
||||
ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts);
|
||||
conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3));
|
||||
conf = ggml_reshape_2d(ctx0, conf, 1, n_tok);
|
||||
|
||||
// note: broadcast the [1, n_tok] confidences to n_embd-wide rows to be able to reuse `llama_get_embeddings_nextn`
|
||||
conf = ggml_repeat(ctx0, conf, res->t_embd);
|
||||
res->t_h_nextn = conf;
|
||||
ggml_build_forward_expand(g.gf, conf);
|
||||
}
|
||||
|
||||
res->t_logits = out;
|
||||
ggml_build_forward_expand(g.gf, out);
|
||||
}
|
||||
|
||||
// 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
|
||||
@@ -210,6 +315,8 @@ 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);
|
||||
|
||||
ggml_tensor * inp_tokens = inp->tokens;
|
||||
|
||||
ggml_tensor * inpL = ggml_get_rows(ctx0, tok_embd, inp->tokens);
|
||||
cb(inpL, "inp_noise_embd", -1);
|
||||
|
||||
@@ -290,4 +397,9 @@ llama_model_dflash::graph<false>::graph(const llama_model & model, const llm_gra
|
||||
res->t_logits = cur;
|
||||
|
||||
ggml_build_forward_expand(gf, cur);
|
||||
|
||||
// DSpark: bias the draft logits with the Markov head
|
||||
if (model.dspark_markov_w1) {
|
||||
build_dspark_markov_head(*this, model, inp_tokens);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user