From 93a1bd94c5edb966d338fb13d247b1bc919971cc Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Thu, 2 Jul 2026 15:24:03 +0800 Subject: [PATCH 01/12] 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). --- common/common.h | 3 +- common/speculative.cpp | 113 ++++++++++++++++++++++++++++- conversion/__init__.py | 1 + conversion/qwen.py | 50 +++++++++++++ gguf-py/gguf/constants.py | 28 ++++++++ gguf-py/gguf/gguf_writer.py | 3 + gguf-py/gguf/tensor_mapping.py | 12 ++++ src/llama-arch.cpp | 10 +++ src/llama-arch.h | 6 ++ src/llama-context.cpp | 2 +- src/llama-hparams.h | 3 + src/llama-model.cpp | 6 +- src/llama-model.h | 6 ++ src/models/dspark.cpp | 126 +++++++++++++++++++++++++++++++++ src/models/models.h | 16 +++++ tests/test-llama-archs.cpp | 4 +- 16 files changed, 381 insertions(+), 8 deletions(-) create mode 100644 src/models/dspark.cpp diff --git a/common/common.h b/common/common.h index bffc1767a7d1..ca02c2964198 100644 --- a/common/common.h +++ b/common/common.h @@ -172,6 +172,7 @@ enum common_speculative_type { COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, // Eagle3 speculative decoding COMMON_SPECULATIVE_TYPE_DRAFT_MTP, // Multi-token prediction COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, // DFlash speculative decoding + COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, // DSpark speculative decoding (DFlash + Markov head) COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding based on n-grams COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values @@ -387,7 +388,7 @@ struct common_params_speculative { uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { - return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH; + return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; }); return needs_rs_seq ? draft.n_max : 0u; diff --git a/common/speculative.cpp b/common/speculative.cpp index 3cb08767bd46..43418eb7ed05 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -34,6 +34,7 @@ const std::map 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}, @@ -924,8 +925,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // scratch buffer for concatenated target features [n_tokens, n_embd_enc] std::vector 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) { auto * ctx_tgt = this->params.ctx_tgt; @@ -1201,6 +1203,101 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } }; +// DSpark: DFlash backbone + a semi-autoregressive Markov head; reuses process(), only draft() differs +struct common_speculative_impl_draft_dspark : public common_speculative_impl_draft_dflash { + common_speculative_impl_draft_dspark(const common_params_speculative & params_in, uint32_t n_seq) + : common_speculative_impl_draft_dflash(params_in, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) + { + auto * ctx_dft = params.ctx_dft; + const llama_model * model_dft = llama_get_model(ctx_dft); + + { + char buf[32] = {}; + if (llama_model_meta_val_str(model_dft, "dspark.block_size", buf, sizeof(buf)) < 0) { + GGML_ABORT("DSpark: missing required metadata key 'dspark.block_size'"); + } + block_size = std::atoi(buf); + } + if (params.n_max > block_size) { + params.n_max = block_size; + } + + LOG_INF("%s: adding speculative implementation 'draft-dspark'\n", __func__); + LOG_INF("%s: - block_size=%d, n_max=%d\n", __func__, block_size, params.n_max); + } + + void draft(common_speculative_draft_params_vec & dparams) override { + auto * ctx_dft = params.ctx_dft; + + common_batch_clear(batch); + + std::vector i_block_beg(n_seq, -1); + std::vector n_block (n_seq, 0); + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + auto & dp = dparams[seq_id]; + if (!dp.drafting) { + continue; + } + + common_sampler_reset(smpls[seq_id].get()); + + 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); + } + n_draft = std::min(n_draft, block_size); + if (n_draft <= 0) { + continue; + } + + // anchor-first block [id_last, * (block_size-1)]: submit the whole block so the + // in-graph Markov head can key anchors off the block boundaries; keep the first n_draft + i_block_beg[seq_id] = batch.n_tokens; + n_block [seq_id] = n_draft; + for (int32_t i = 0; i < block_size; ++i) { + common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); + } + } + + if (batch.n_tokens == 0) { + return; + } + + if (llama_decode(ctx_dft, batch) != 0) { + LOG_WRN("%s: llama_decode failed\n", __func__); + return; + } + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + if (i_block_beg[seq_id] < 0) { + continue; + } + auto & dp = dparams[seq_id]; + auto & result = *dp.result; + + const int32_t beg = i_block_beg[seq_id]; + const int32_t nb = n_block[seq_id]; // drafts to keep (<= block_size) + + auto * smpl = smpls[seq_id].get(); + // greedily read the predicted block at this sequence's noise positions 1..nb-1 + for (int32_t i = 0; i < nb; ++i) { + common_sampler_sample(smpl, ctx_dft, beg + i, true); + + const auto * cur_p = common_sampler_get_candidates(smpl, true); + + const llama_token id = cur_p->data[0].id; + + common_sampler_accept(smpl, id, true); + + result.push_back(id); + } + } + } +}; + struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) @@ -2145,6 +2242,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"; @@ -2198,6 +2296,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: @@ -2342,6 +2441,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; @@ -2352,7 +2452,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 @@ -2385,6 +2485,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> impls = {}; @@ -2409,6 +2512,10 @@ common_speculative * common_speculative_init(common_params_speculative & params, impls.push_back(std::make_unique(config.params, n_seq)); break; } + case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: { + impls.push_back(std::make_unique(config.params, n_seq)); + break; + } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { common_ngram_map ngram_map = get_common_ngram_map(config.type, config.params.ngram_simple); diff --git a/conversion/__init__.py b/conversion/__init__.py index 465c010f667a..8575867e77e7 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -52,6 +52,7 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "Qwen3DSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", diff --git a/conversion/qwen.py b/conversion/qwen.py index 9bc2b99fde5d..0a1f99076b1f 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -688,3 +688,53 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca if not name.startswith("model."): name = "model." + name return super().filter_tensors((name, gen)) + + +@ModelBase.register("Qwen3DSparkModel") +class DSparkModel(Qwen3Model): + # DSpark = DFlash backbone + a semi-autoregressive Markov head (+ optional confidence head). + # The DeepSpec checkpoint stores its config flat (block_size / target_layer_ids / mask_token_id / + # markov_rank at top level). embed_tokens / lm_head are byte-identical to the target, so they are + # NOT emitted here -- the DSpark decoder shares the target's via ctx_other (same as DFlash). + model_arch = gguf.MODEL_ARCH.DSPARK + + def set_vocab(self): + if self.target_model_dir is None: + raise ValueError( + "DSpark draft model requires --target-model-dir to be specified. " + "Please provide the path to the target model directory containing the tokenizer." + ) + logger.info(f"DSpark: Using tokenizer from target model: {self.target_model_dir}") + original_dir = self.dir_model + self.dir_model = self.target_model_dir + super().set_vocab() + self.dir_model = original_dir + + mask_token_id = self.hparams.get("mask_token_id") + if mask_token_id is not None: + self.gguf_writer.add_mask_token_id(mask_token_id) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + block_size = self.hparams.get("block_size", 7) + self.gguf_writer.add_block_size(block_size) + + # flat DeepSpec schema; mirror DFlash's +1 extract-layer convention + target_layer_ids = self.hparams.get("target_layer_ids", []) + if target_layer_ids: + extract_layer_ids = [i + 1 for i in target_layer_ids] + self.gguf_writer.add_target_layers(extract_layer_ids) + + markov_rank = self.hparams.get("markov_rank", 0) + self.gguf_writer.add_markov_rank(markov_rank) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + # embed_tokens / lm_head are byte-identical to the target and shared at runtime -- drop them + if name.endswith(("embed_tokens.weight", "lm_head.weight")): + return None + if not name.startswith("model."): + name = "model." + name + return super().filter_tensors((name, gen)) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 63ac2ed1f5cf..4528d2faead1 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -158,6 +158,7 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + MARKOV_RANK = "{arch}.markov_rank" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" class Attention: @@ -531,6 +532,7 @@ class MODEL_ARCH(IntEnum): MISTRAL3 = auto() EAGLE3 = auto() DFLASH = auto() + DSPARK = auto() MISTRAL4 = auto() PADDLEOCR = auto() MIMO2 = auto() @@ -954,6 +956,9 @@ class MODEL_TENSOR(IntEnum): # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping + DSPARK_MARKOV_W1 = auto() # dspark markov head: prev-token embed + DSPARK_MARKOV_W2 = auto() # dspark markov head: bias projection + DSPARK_CONF_PROJ = auto() # dspark confidence head: proj # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1113,6 +1118,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.MISTRAL3: "mistral3", MODEL_ARCH.EAGLE3: "eagle3", MODEL_ARCH.DFLASH: "dflash", + MODEL_ARCH.DSPARK: "dspark", MODEL_ARCH.MISTRAL4: "mistral4", MODEL_ARCH.PADDLEOCR: "paddleocr", MODEL_ARCH.MIMO2: "mimo2", @@ -1561,6 +1567,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", MODEL_TENSOR.FC: "fc", + MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", + MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", + MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", MODEL_TENSOR.D2T: "d2t", } @@ -4237,6 +4246,25 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, ], + MODEL_ARCH.DSPARK: [ + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FC, + MODEL_TENSOR.ENC_OUTPUT_NORM, + MODEL_TENSOR.DSPARK_MARKOV_W1, + MODEL_TENSOR.DSPARK_MARKOV_W2, + MODEL_TENSOR.DSPARK_CONF_PROJ, + ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 1e277f0687c5..09522d79aab9 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -946,6 +946,9 @@ def add_sliding_window(self, value: int) -> None: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) + def add_markov_rank(self, value: int) -> None: + self.add_uint32(Keys.LLM.MARKOV_RANK.format(arch=self.arch), value) + def add_target_layers(self, value: Sequence[int]) -> None: self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 9efb36f8a447..011dae886789 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1290,6 +1290,18 @@ class TensorNameMap: "model.fc", # dflash ), + MODEL_TENSOR.DSPARK_MARKOV_W1: ( + "model.markov_head.markov_w1", # dspark + ), + + MODEL_TENSOR.DSPARK_MARKOV_W2: ( + "model.markov_head.markov_w2", # dspark + ), + + MODEL_TENSOR.DSPARK_CONF_PROJ: ( + "model.confidence_head.proj", # dspark + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 72968607db80..14ea241249b3 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -132,6 +132,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MISTRAL3, "mistral3" }, { LLM_ARCH_EAGLE3, "eagle3" }, { LLM_ARCH_DFLASH, "dflash" }, + { LLM_ARCH_DSPARK, "dspark" }, { LLM_ARCH_MISTRAL4, "mistral4" }, { LLM_ARCH_PADDLEOCR, "paddleocr" }, { LLM_ARCH_MIMO2, "mimo2" }, @@ -308,6 +309,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TARGET_LAYERS, "%s.target_layers" }, { LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" }, + { LLM_KV_MARKOV_RANK, "%s.markov_rank" }, + { LLM_KV_BLOCK_SIZE, "%s.block_size" }, { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, @@ -604,6 +607,9 @@ static const std::map 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: @@ -855,6 +861,10 @@ static const std::map 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) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index b74d53af4a32..ce4bde8a01ec 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -146,6 +146,7 @@ enum llm_arch { LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, LLM_ARCH_DFLASH, + LLM_ARCH_DSPARK, LLM_ARCH_UNKNOWN, }; @@ -354,6 +355,8 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, + LLM_KV_MARKOV_RANK, + LLM_KV_BLOCK_SIZE, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_SHORTCONV_L_CACHE, @@ -612,6 +615,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, }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index eed041eef4e7..35f98f06e28d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -149,7 +149,7 @@ llama_context::llama_context( cparams.ctx_other = params.ctx_other; } - if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) { + if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH || model.arch == LLM_ARCH_DSPARK) { if (model.tok_embd == nullptr || model.output == nullptr) { if (params.ctx_other == nullptr) { throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)"); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 8be5f28f39e6..e0e4ef245bfd 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -187,6 +187,9 @@ struct llama_hparams { // for Classifiers uint32_t n_cls_out = 1; + // for DSpark: the trained draft block size, in tokens (anchor + n-1 masks) + uint32_t n_dspark_block = 0; + // input embedding dimension (0 = use n_embd) uint32_t n_embd_inp_impl = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4c10e4126f62..d4798465b1be 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -298,6 +298,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_eagle3(params); case LLM_ARCH_DFLASH: return new llama_model_dflash(params); + case LLM_ARCH_DSPARK: + return new llama_model_dspark(params); case LLM_ARCH_MIMO2: return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: @@ -2555,6 +2557,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_TALKIE: case LLM_ARCH_MELLUM: case LLM_ARCH_DFLASH: + case LLM_ARCH_DSPARK: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_QWEN2VL: @@ -2683,7 +2686,8 @@ bool llama_model_has_encoder(const llama_model * model) { case LLM_ARCH_T5: case LLM_ARCH_T5ENCODER: case LLM_ARCH_EAGLE3: - case LLM_ARCH_DFLASH: return true; + case LLM_ARCH_DFLASH: + case LLM_ARCH_DSPARK: return true; default: return false; } } diff --git a/src/llama-model.h b/src/llama-model.h index 45b054cedf1d..35ab5f0ff8b4 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -599,6 +599,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 target_layer_ids; diff --git a/src/models/dspark.cpp b/src/models/dspark.cpp new file mode 100644 index 000000000000..1c873407ab4a --- /dev/null +++ b/src/models/dspark.cpp @@ -0,0 +1,126 @@ +#include "models.h" + +// DSpark = DFlash backbone + a semi-autoregressive Markov head applied in-graph by the decoder + +void llama_model_dspark::load_arch_hparams(llama_model_loader & ml) { + llama_model_dflash::load_arch_hparams(ml); + + ml.get_key(LLM_KV_BLOCK_SIZE, hparams.n_dspark_block, /*required*/ true); +} + +void llama_model_dspark::load_arch_tensors(llama_model_loader & ml) { + llama_model_dflash::load_arch_tensors(ml); + + LLAMA_LOAD_LOCALS; + + uint32_t markov_rank = 0; + ml.get_key(LLM_KV_MARKOV_RANK, markov_rank, /*required*/ true); + const int64_t R = (int64_t) markov_rank; + + dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { R, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { R, n_vocab }, 0); + + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + R, 1 }, TENSOR_NOT_REQUIRED); + dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); +} + +std::unique_ptr llama_model_dspark::build_arch_graph(const llm_graph_params & params) const { + switch (params.gtype) { + case LLM_GRAPH_TYPE_ENCODER: + return std::make_unique>(*this, params); + case LLM_GRAPH_TYPE_DEFAULT: + case LLM_GRAPH_TYPE_DECODER: + return std::make_unique>(*this, params); + default: + GGML_ABORT("invalid graph type"); + }; +} + +// DSpark encoder == DFlash encoder +template <> +llama_model_dspark::graph::graph(const llama_model & model, const llm_graph_params & params) + : llama_model_dflash::graph(model, params) {} + +// anchor (committed last) token of every draft block: token 0 of each block in the ubatch +class llm_graph_input_dspark_anchor : public llm_graph_input_i { +public: + llm_graph_input_dspark_anchor(uint32_t block_size) : block_size(block_size) {} + virtual ~llm_graph_input_dspark_anchor() = default; + + void set_input(const llama_ubatch * ubatch) override { + GGML_ASSERT(ubatch->token); + const int64_t n_blocks = anchors->ne[0]; + std::vector buf(n_blocks); + for (int64_t j = 0; j < n_blocks; ++j) { + buf[j] = ubatch->token[j*block_size]; + } + ggml_backend_tensor_set(anchors, buf.data(), 0, n_blocks*sizeof(int32_t)); + } + + bool can_reuse(const llm_graph_params & params) override { + return params.ubatch.token && anchors && + anchors->ne[0]*(int64_t) block_size == (int64_t) params.ubatch.n_tokens; + } + + ggml_tensor * anchors = nullptr; // I32 [n_blocks] + + const uint32_t block_size; +}; + +// DSpark decoder: DFlash decoder + Markov bias on the draft logits, 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)) +template <> +llama_model_dspark::graph::graph(const llama_model & model, const llm_graph_params & params) + : llama_model_dflash::graph(model, params) { + // KV-injection (embd) batch: no logits to bias + if (ubatch.embd) { + return; + } + + ggml_tensor * w1 = model.dspark_markov_w1; + ggml_tensor * w2 = model.dspark_markov_w2; + GGML_ASSERT(w1 && w2 && "DSpark markov 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 int64_t bs = model.hparams.n_dspark_block; + GGML_ASSERT(bs > 0); + + // the drafting loop always submits whole anchor-first blocks + if (n_tok % bs != 0) { + return; + } + const int64_t n_blocks = n_tok / bs; + + auto inp = std::make_unique((uint32_t) bs); + inp->anchors = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_blocks); + ggml_set_input(inp->anchors); + ggml_tensor * prev = inp->anchors; // I32 [n_blocks] + res->add_input(std::move(inp)); + + ggml_tensor * cat = nullptr; + for (int64_t i = 0; i < bs; ++i) { + ggml_tensor * bias = ggml_mul_mat(ctx0, w2, ggml_get_rows(ctx0, 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, bs*base->nb[1], i*base->nb[1]); + ggml_tensor * col = ggml_add(ctx0, base_i, bias); + + cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; + + if (i + 1 < bs) { + prev = ggml_argmax(ctx0, col); // I32 [n_blocks] + } + } + + // cat is position-major; restore the ubatch's block-major order + ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, bs); + out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs, n_blocks] + out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); + + res->t_logits = out; + ggml_build_forward_expand(gf, out); +} diff --git a/src/models/models.h b/src/models/models.h index a86ae05aae6b..034e1282e338 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1254,6 +1254,22 @@ struct llama_model_dflash : public llama_model_base { }; +struct llama_model_dspark : public llama_model_dflash { + llama_model_dspark(const struct llama_model_params & params) : llama_model_dflash(params) {} + // extend the DFlash hparams/tensors with the block size and the Markov / confidence heads + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + // the DFlash graphs plus the in-graph Markov head on the decoder's draft logits + template + struct graph : public llama_model_dflash::graph { + graph(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_mistral4 : public llama_model_deepseek2 { llama_model_mistral4(const struct llama_model_params & params) : llama_model_deepseek2(params) {} // reuse load_arch_hparams and load_arch_tensors from llama_model_deepseek2 diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 57d33a627d5d..fb71e9b0ec9c 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -457,7 +457,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } - if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { + if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH || arch == LLM_ARCH_DSPARK) { continue; } for (bool moe : {false, true}) { @@ -563,7 +563,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } - if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { + if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH || arch == LLM_ARCH_DSPARK) { continue; } From 61d2ee5d6869cec3a7d006993756c735d3ce015b Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Sat, 4 Jul 2026 14:44:41 +0800 Subject: [PATCH 02/12] spec: read draft block size in the dflash impl --- common/speculative.cpp | 48 ++++++++++++++++-------------------------- 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 43418eb7ed05..11760bbb62e7 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -926,7 +926,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::vector features_buf; 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_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, + bool draft_at_anchor = false) : common_speculative_impl(type, n_seq) , params(params.draft) { @@ -945,26 +946,31 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { n_embd_dec = llama_model_n_embd(model_dft); n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; - // read the trained block size from the dflash.block_size metadata key + // read the trained block size from the .block_size metadata key block_size = 16; { + char arch[64] = {}; + llama_model_meta_val_str(model_dft, "general.architecture", arch, sizeof(arch)); + char buf[32] = {}; - if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) { + if (llama_model_meta_val_str(model_dft, (std::string(arch) + ".block_size").c_str(), buf, sizeof(buf)) >= 0) { block_size = std::atoi(buf); } } 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, * (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); + // the input block is [id_last, * (block_size-1)], so a step drafts at most block_size-1 + // tokens from the mask positions, plus one more when the head also drafts at the anchor position + const int32_t n_draft_max = draft_at_anchor ? 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); @@ -1205,26 +1211,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // DSpark: DFlash backbone + a semi-autoregressive Markov head; reuses process(), only draft() differs struct common_speculative_impl_draft_dspark : public common_speculative_impl_draft_dflash { - common_speculative_impl_draft_dspark(const common_params_speculative & params_in, uint32_t n_seq) - : common_speculative_impl_draft_dflash(params_in, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) - { - auto * ctx_dft = params.ctx_dft; - const llama_model * model_dft = llama_get_model(ctx_dft); - - { - char buf[32] = {}; - if (llama_model_meta_val_str(model_dft, "dspark.block_size", buf, sizeof(buf)) < 0) { - GGML_ABORT("DSpark: missing required metadata key 'dspark.block_size'"); - } - block_size = std::atoi(buf); - } - if (params.n_max > block_size) { - params.n_max = block_size; - } - - LOG_INF("%s: adding speculative implementation 'draft-dspark'\n", __func__); - LOG_INF("%s: - block_size=%d, n_max=%d\n", __func__, block_size, params.n_max); - } + common_speculative_impl_draft_dspark(const common_params_speculative & params, uint32_t n_seq) + : common_speculative_impl_draft_dflash(params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, /*draft_at_anchor*/ true) {} void draft(common_speculative_draft_params_vec & dparams) override { auto * ctx_dft = params.ctx_dft; From e492dd32a6381791783901dfb1d0e143e4dc72b1 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Sat, 4 Jul 2026 14:44:53 +0800 Subject: [PATCH 03/12] docs: add DSpark section to speculative.md --- docs/speculative.md | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/docs/speculative.md b/docs/speculative.md index 4100b92f8f18..810e05e60900 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -78,6 +78,32 @@ See: - #22105 +### DSpark (`draft-dspark`) + +DSpark extends DFlash with a semi-autoregressive _Markov head_: the draft still emits a whole +block per forward pass, but each block position's logits are biased by a low-rank term keyed on +the previous token, chained in-graph across the block. This keeps drafting at one decode per +block while recovering some of the left-to-right signal that pure block diffusion loses. + +The draft is a small DeepSpec checkpoint trained for a specific target (for example +[`deepseek-ai/dspark_qwen3_4b_block7`](https://huggingface.co/deepseek-ai/dspark_qwen3_4b_block7) +for `Qwen/Qwen3-4B`). Convert it with `--target-model-dir` so it inherits the target's tokenizer +and token embeddings: + +```bash +python convert_hf_to_gguf.py deepseek-ai/dspark_qwen3_4b_block7 \ + --target-model-dir Qwen/Qwen3-4B --outtype bf16 --outfile Qwen3-4B-DSpark.gguf + +llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ + --spec-type draft-dspark --spec-draft-n-max 7 -fa on --jinja +``` + +`--spec-draft-n-max` is clamped to the draft model's trained block size. + +See: + +- #25173 + ### n-gram Cache (`ngram-cache`) An n-gram is a sequence of n tokens. The n-gram cache implementation maintains statistics about short n-gram sequences. @@ -173,7 +199,7 @@ If a draft model is combined with a draftless decoding the draftless decoding ha ### General Speculative Parameters ``` ---spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] +--spec-type [none|draft-simple|draft-eagle3|draft-dflash|draft-dspark|draft-mtp|ngram-cache|ngram-simple|ngram-map-k|ngram-map-k4v|ngram-mod] comma-separated list of types of speculative decoding to use (default: none) (env: LLAMA_ARG_SPEC_TYPE) @@ -314,6 +340,7 @@ Specifies a comma-separated list of speculative decoding types to use. | `draft-simple` | Use a simple draft model for speculation | | `draft-eagle3` | Use an EAGLE-3 draft model that reads the target's hidden states | | `draft-dflash` | Use a DFlash block-diffusion draft model that emits a block per step | +| `draft-dspark` | Use a DSpark draft model (DFlash backbone + semi-autoregressive Markov head) | | `draft-mtp` | Use Multi Token Prediction (MTP) heads from the main model | | `ngram-cache` | Use n-gram cache lookup | | `ngram-simple` | Use simple n-gram pattern matching | From f553f76d31a7e45b7d508df77783e42d36925f08 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Sat, 4 Jul 2026 14:59:45 +0800 Subject: [PATCH 04/12] spec: keep dspark block size read in the dspark impl --- common/speculative.cpp | 48 ++++++++++++++++++++++++++---------------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 11760bbb62e7..43418eb7ed05 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -926,8 +926,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::vector features_buf; common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq, - common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH, - bool draft_at_anchor = false) + common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) : common_speculative_impl(type, n_seq) , params(params.draft) { @@ -946,31 +945,26 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { n_embd_dec = llama_model_n_embd(model_dft); n_embd_enc = (int32_t) target_layer_ids_n * n_embd_tgt; - // read the trained block size from the .block_size metadata key + // read the trained block size from the dflash.block_size metadata key block_size = 16; { - char arch[64] = {}; - llama_model_meta_val_str(model_dft, "general.architecture", arch, sizeof(arch)); - char buf[32] = {}; - if (llama_model_meta_val_str(model_dft, (std::string(arch) + ".block_size").c_str(), buf, sizeof(buf)) >= 0) { + if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) { block_size = std::atoi(buf); } } mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); - LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); + LOG_INF("%s: adding speculative implementation 'draft-dflash'\n", __func__); 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); - // the input block is [id_last, * (block_size-1)], so a step drafts at most block_size-1 - // tokens from the mask positions, plus one more when the head also drafts at the anchor position - const int32_t n_draft_max = draft_at_anchor ? 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); + // DFlash input is [id_last, * (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); } batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); @@ -1211,8 +1205,26 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // DSpark: DFlash backbone + a semi-autoregressive Markov head; reuses process(), only draft() differs struct common_speculative_impl_draft_dspark : public common_speculative_impl_draft_dflash { - common_speculative_impl_draft_dspark(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl_draft_dflash(params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK, /*draft_at_anchor*/ true) {} + common_speculative_impl_draft_dspark(const common_params_speculative & params_in, uint32_t n_seq) + : common_speculative_impl_draft_dflash(params_in, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) + { + auto * ctx_dft = params.ctx_dft; + const llama_model * model_dft = llama_get_model(ctx_dft); + + { + char buf[32] = {}; + if (llama_model_meta_val_str(model_dft, "dspark.block_size", buf, sizeof(buf)) < 0) { + GGML_ABORT("DSpark: missing required metadata key 'dspark.block_size'"); + } + block_size = std::atoi(buf); + } + if (params.n_max > block_size) { + params.n_max = block_size; + } + + LOG_INF("%s: adding speculative implementation 'draft-dspark'\n", __func__); + LOG_INF("%s: - block_size=%d, n_max=%d\n", __func__, block_size, params.n_max); + } void draft(common_speculative_draft_params_vec & dparams) override { auto * ctx_dft = params.ctx_dft; From a8ba490077692b4c7bca177a04b0c32a622ee920 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Wed, 8 Jul 2026 17:27:28 +0800 Subject: [PATCH 05/12] 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) --- common/speculative.cpp | 3 +++ docs/speculative.md | 3 +++ src/models/dspark.cpp | 8 ++++++++ 3 files changed, 14 insertions(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index 43418eb7ed05..54b7892dbb40 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1283,6 +1283,9 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl_dra auto * smpl = smpls[seq_id].get(); // greedily read the predicted block at this sequence's noise positions 1..nb-1 + // TODO: implement confidence-scheduled prefix pruning (the "scheduled" part of DSpark): + // use the confidence head to truncate the drafted block at the first low-confidence + // position instead of always keeping all n_draft tokens for (int32_t i = 0; i < nb; ++i) { common_sampler_sample(smpl, ctx_dft, beg + i, true); diff --git a/docs/speculative.md b/docs/speculative.md index 810e05e60900..84d45d47c54d 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -100,6 +100,9 @@ llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ `--spec-draft-n-max` is clamped to the draft model's trained block size. +Currently only drafts with a Qwen3 backbone are supported; support for other backbones +(e.g. Gemma4) is planned. + See: - #25173 diff --git a/src/models/dspark.cpp b/src/models/dspark.cpp index 1c873407ab4a..30a8583b601c 100644 --- a/src/models/dspark.cpp +++ b/src/models/dspark.cpp @@ -1,6 +1,9 @@ #include "models.h" // DSpark = DFlash backbone + a semi-autoregressive Markov head applied in-graph by the decoder +// +// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) +// need their own conversion path and graph tweaks in the shared DFlash code void llama_model_dspark::load_arch_hparams(llama_model_loader & ml) { llama_model_dflash::load_arch_hparams(ml); @@ -20,6 +23,9 @@ void llama_model_dspark::load_arch_tensors(llama_model_loader & ml) { dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { R, n_vocab }, 0); dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { R, n_vocab }, 0); + // TODO: the confidence head is loaded but not yet used by the graph -- it should gate + // the draft length per block (early-exit on low confidence) instead of always + // drafting the full block dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + R, 1 }, TENSOR_NOT_REQUIRED); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); } @@ -102,6 +108,8 @@ llama_model_dspark::graph::graph(const llama_model & model, const llm_gra res->add_input(std::move(inp)); ggml_tensor * cat = 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 < bs; ++i) { ggml_tensor * bias = ggml_mul_mat(ctx0, w2, ggml_get_rows(ctx0, w1, prev)); // [n_vocab, n_blocks] From 65930be5f897483bc92eb9cc84424239740b33b1 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Thu, 9 Jul 2026 08:42:59 +0800 Subject: [PATCH 06/12] 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. --- common/speculative.cpp | 20 ++---- conversion/qwen.py | 48 +++---------- gguf-py/gguf/constants.py | 20 +----- gguf-py/gguf/gguf_writer.py | 3 - src/llama-arch.cpp | 3 - src/llama-arch.h | 3 - src/llama-context.cpp | 2 +- src/llama-hparams.h | 3 - src/llama-model.cpp | 6 +- src/models/dflash.cpp | 85 +++++++++++++++++++++++ src/models/dspark.cpp | 134 ------------------------------------ src/models/models.h | 16 ----- tests/test-llama-archs.cpp | 4 +- 13 files changed, 105 insertions(+), 242 deletions(-) delete mode 100644 src/models/dspark.cpp diff --git a/common/speculative.cpp b/common/speculative.cpp index 54b7892dbb40..2b9836ba188a 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1203,24 +1203,16 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } }; -// DSpark: DFlash backbone + a semi-autoregressive Markov head; reuses process(), only draft() differs +// DSpark: DFlash backbone + a semi-autoregressive Markov head; reuses process(), only draft() differs. +// The draft model is a DFlash GGUF carrying the Markov head tensors -- the decoder graph picks them +// up by presence, this impl only changes how the noise blocks are submitted and read back. struct common_speculative_impl_draft_dspark : public common_speculative_impl_draft_dflash { common_speculative_impl_draft_dspark(const common_params_speculative & params_in, uint32_t n_seq) : common_speculative_impl_draft_dflash(params_in, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { - auto * ctx_dft = params.ctx_dft; - const llama_model * model_dft = llama_get_model(ctx_dft); - - { - char buf[32] = {}; - if (llama_model_meta_val_str(model_dft, "dspark.block_size", buf, sizeof(buf)) < 0) { - GGML_ABORT("DSpark: missing required metadata key 'dspark.block_size'"); - } - block_size = std::atoi(buf); - } - if (params.n_max > block_size) { - params.n_max = block_size; - } + // DSpark drafts from position 0 of the block (the anchor position), so unlike DFlash + // it can draft a full block_size tokens per step; undo the base class's n_max clamp + params.n_max = std::min(params_in.draft.n_max, block_size); LOG_INF("%s: adding speculative implementation 'draft-dspark'\n", __func__); LOG_INF("%s: - block_size=%d, n_max=%d\n", __func__, block_size, params.n_max); diff --git a/conversion/qwen.py b/conversion/qwen.py index 0a1f99076b1f..d1127f7431f6 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -691,50 +691,20 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca @ModelBase.register("Qwen3DSparkModel") -class DSparkModel(Qwen3Model): - # DSpark = DFlash backbone + a semi-autoregressive Markov head (+ optional confidence head). - # The DeepSpec checkpoint stores its config flat (block_size / target_layer_ids / mask_token_id / - # markov_rank at top level). embed_tokens / lm_head are byte-identical to the target, so they are - # NOT emitted here -- the DSpark decoder shares the target's via ctx_other (same as DFlash). - model_arch = gguf.MODEL_ARCH.DSPARK - - def set_vocab(self): - if self.target_model_dir is None: - raise ValueError( - "DSpark draft model requires --target-model-dir to be specified. " - "Please provide the path to the target model directory containing the tokenizer." - ) - logger.info(f"DSpark: Using tokenizer from target model: {self.target_model_dir}") - original_dir = self.dir_model - self.dir_model = self.target_model_dir - super().set_vocab() - self.dir_model = original_dir - - mask_token_id = self.hparams.get("mask_token_id") - if mask_token_id is not None: - self.gguf_writer.add_mask_token_id(mask_token_id) - - def set_gguf_parameters(self): - super().set_gguf_parameters() - - block_size = self.hparams.get("block_size", 7) - self.gguf_writer.add_block_size(block_size) - - # flat DeepSpec schema; mirror DFlash's +1 extract-layer convention - target_layer_ids = self.hparams.get("target_layer_ids", []) - if target_layer_ids: - extract_layer_ids = [i + 1 for i in target_layer_ids] - self.gguf_writer.add_target_layers(extract_layer_ids) +class DSparkModel(DFlashModel): + # DSpark = DFlash + a semi-autoregressive Markov head + model_arch = gguf.MODEL_ARCH.DFLASH - markov_rank = self.hparams.get("markov_rank", 0) - self.gguf_writer.add_markov_rank(markov_rank) + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # normalize the flat DeepSpec schema to DFlash's nested dflash_config + self.hparams.setdefault("dflash_config", { + k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams + }) @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item - # embed_tokens / lm_head are byte-identical to the target and shared at runtime -- drop them if name.endswith(("embed_tokens.weight", "lm_head.weight")): return None - if not name.startswith("model."): - name = "model." + name return super().filter_tensors((name, gen)) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 4528d2faead1..c1895c0fe0b8 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -158,7 +158,6 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" - MARKOV_RANK = "{arch}.markov_rank" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" class Attention: @@ -532,7 +531,6 @@ class MODEL_ARCH(IntEnum): MISTRAL3 = auto() EAGLE3 = auto() DFLASH = auto() - DSPARK = auto() MISTRAL4 = auto() PADDLEOCR = auto() MIMO2 = auto() @@ -1118,7 +1116,6 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.MISTRAL3: "mistral3", MODEL_ARCH.EAGLE3: "eagle3", MODEL_ARCH.DFLASH: "dflash", - MODEL_ARCH.DSPARK: "dspark", MODEL_ARCH.MISTRAL4: "mistral4", MODEL_ARCH.PADDLEOCR: "paddleocr", MODEL_ARCH.MIMO2: "mimo2", @@ -4245,22 +4242,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, - ], - MODEL_ARCH.DSPARK: [ - MODEL_TENSOR.OUTPUT_NORM, - MODEL_TENSOR.ATTN_NORM, - MODEL_TENSOR.ATTN_Q, - MODEL_TENSOR.ATTN_K, - MODEL_TENSOR.ATTN_V, - MODEL_TENSOR.ATTN_OUT, - MODEL_TENSOR.ATTN_Q_NORM, - MODEL_TENSOR.ATTN_K_NORM, - MODEL_TENSOR.FFN_NORM, - MODEL_TENSOR.FFN_GATE, - MODEL_TENSOR.FFN_DOWN, - MODEL_TENSOR.FFN_UP, - MODEL_TENSOR.FC, - MODEL_TENSOR.ENC_OUTPUT_NORM, + # optional DSpark heads (present only in DSpark drafts) MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 09522d79aab9..1e277f0687c5 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -946,9 +946,6 @@ def add_sliding_window(self, value: int) -> None: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) - def add_markov_rank(self, value: int) -> None: - self.add_uint32(Keys.LLM.MARKOV_RANK.format(arch=self.arch), value) - def add_target_layers(self, value: Sequence[int]) -> None: self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 14ea241249b3..5fc840ba8925 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -132,7 +132,6 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MISTRAL3, "mistral3" }, { LLM_ARCH_EAGLE3, "eagle3" }, { LLM_ARCH_DFLASH, "dflash" }, - { LLM_ARCH_DSPARK, "dspark" }, { LLM_ARCH_MISTRAL4, "mistral4" }, { LLM_ARCH_PADDLEOCR, "paddleocr" }, { LLM_ARCH_MIMO2, "mimo2" }, @@ -309,8 +308,6 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_TARGET_LAYERS, "%s.target_layers" }, { LLM_KV_TARGET_HIDDEN_SIZE, "%s.target_hidden_size" }, - { LLM_KV_MARKOV_RANK, "%s.markov_rank" }, - { LLM_KV_BLOCK_SIZE, "%s.block_size" }, { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index ce4bde8a01ec..cf5af7033bd7 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -146,7 +146,6 @@ enum llm_arch { LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, LLM_ARCH_DFLASH, - LLM_ARCH_DSPARK, LLM_ARCH_UNKNOWN, }; @@ -355,8 +354,6 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, - LLM_KV_MARKOV_RANK, - LLM_KV_BLOCK_SIZE, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_SHORTCONV_L_CACHE, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 35f98f06e28d..eed041eef4e7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -149,7 +149,7 @@ llama_context::llama_context( cparams.ctx_other = params.ctx_other; } - if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH || model.arch == LLM_ARCH_DSPARK) { + if (model.arch == LLM_ARCH_EAGLE3 || model.arch == LLM_ARCH_DFLASH) { if (model.tok_embd == nullptr || model.output == nullptr) { if (params.ctx_other == nullptr) { throw std::runtime_error(model.arch_name() + " requires ctx_other to be set (this warning is normal during memory fitting)"); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index e0e4ef245bfd..8be5f28f39e6 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -187,9 +187,6 @@ struct llama_hparams { // for Classifiers uint32_t n_cls_out = 1; - // for DSpark: the trained draft block size, in tokens (anchor + n-1 masks) - uint32_t n_dspark_block = 0; - // input embedding dimension (0 = use n_embd) uint32_t n_embd_inp_impl = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d4798465b1be..4c10e4126f62 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -298,8 +298,6 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_eagle3(params); case LLM_ARCH_DFLASH: return new llama_model_dflash(params); - case LLM_ARCH_DSPARK: - return new llama_model_dspark(params); case LLM_ARCH_MIMO2: return new llama_model_mimo2(params); case LLM_ARCH_KIMI_LINEAR: @@ -2557,7 +2555,6 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_TALKIE: case LLM_ARCH_MELLUM: case LLM_ARCH_DFLASH: - case LLM_ARCH_DSPARK: return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_QWEN2VL: @@ -2686,8 +2683,7 @@ bool llama_model_has_encoder(const llama_model * model) { case LLM_ARCH_T5: case LLM_ARCH_T5ENCODER: case LLM_ARCH_EAGLE3: - case LLM_ARCH_DFLASH: - case LLM_ARCH_DSPARK: return true; + case LLM_ARCH_DFLASH: return true; default: return false; } } diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index a7b4f4435a88..7894e9dfa49a 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -36,6 +36,26 @@ 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; detect it by tensor presence + // + // 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 R = markov_meta->ne[0]; // markov rank + + dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { R, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { R, n_vocab }, 0); + + // TODO: the confidence head is loaded but not yet used by the graph -- it should gate + // the draft length per block (early-exit on low confidence) instead of always + // drafting the full block + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + R, 1 }, TENSOR_NOT_REQUIRED); + 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) R); + } + 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 @@ -104,6 +124,64 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_grap ggml_build_forward_expand(gf, cur); } +// DSpark drafts only (DFlash + Markov head): Markov bias on the draft logits, +// 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)) +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 && "DSpark markov 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]; + + // the trained draft block size, in tokens (anchor + n-1 masks) + 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 bs = std::stoi(it->second); + GGML_ASSERT(bs > 0); + + // the drafting loop always submits whole anchor-first blocks + if (n_tok % bs != 0) { + return; + } + const int64_t n_blocks = n_tok / bs; + + // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view + ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, bs*tokens->nb[0], 0); + prev = ggml_cont_1d(ctx0, prev, n_blocks); // I32 [n_blocks] + + ggml_tensor * cat = 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 < bs; ++i) { + ggml_tensor * bias = ggml_mul_mat(ctx0, w2, ggml_get_rows(ctx0, 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, bs*base->nb[1], i*base->nb[1]); + ggml_tensor * col = ggml_add(ctx0, base_i, bias); + + cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; + + if (i + 1 < bs) { + prev = ggml_argmax(ctx0, col); // I32 [n_blocks] + } + } + + // cat is position-major; restore the ubatch's block-major order + ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, bs); + out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs, n_blocks] + out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); + + 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 @@ -193,6 +271,8 @@ llama_model_dflash::graph::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); @@ -273,4 +353,9 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra res->t_logits = cur; ggml_build_forward_expand(gf, cur); + + // DSpark drafts only: bias the draft logits with the Markov head + if (model.dspark_markov_w1) { + build_dspark_markov_head(*this, model, inp_tokens); + } } diff --git a/src/models/dspark.cpp b/src/models/dspark.cpp deleted file mode 100644 index 30a8583b601c..000000000000 --- a/src/models/dspark.cpp +++ /dev/null @@ -1,134 +0,0 @@ -#include "models.h" - -// DSpark = DFlash backbone + a semi-autoregressive Markov head applied in-graph by the decoder -// -// TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) -// need their own conversion path and graph tweaks in the shared DFlash code - -void llama_model_dspark::load_arch_hparams(llama_model_loader & ml) { - llama_model_dflash::load_arch_hparams(ml); - - ml.get_key(LLM_KV_BLOCK_SIZE, hparams.n_dspark_block, /*required*/ true); -} - -void llama_model_dspark::load_arch_tensors(llama_model_loader & ml) { - llama_model_dflash::load_arch_tensors(ml); - - LLAMA_LOAD_LOCALS; - - uint32_t markov_rank = 0; - ml.get_key(LLM_KV_MARKOV_RANK, markov_rank, /*required*/ true); - const int64_t R = (int64_t) markov_rank; - - dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { R, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { R, n_vocab }, 0); - - // TODO: the confidence head is loaded but not yet used by the graph -- it should gate - // the draft length per block (early-exit on low confidence) instead of always - // drafting the full block - dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + R, 1 }, TENSOR_NOT_REQUIRED); - dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); -} - -std::unique_ptr llama_model_dspark::build_arch_graph(const llm_graph_params & params) const { - switch (params.gtype) { - case LLM_GRAPH_TYPE_ENCODER: - return std::make_unique>(*this, params); - case LLM_GRAPH_TYPE_DEFAULT: - case LLM_GRAPH_TYPE_DECODER: - return std::make_unique>(*this, params); - default: - GGML_ABORT("invalid graph type"); - }; -} - -// DSpark encoder == DFlash encoder -template <> -llama_model_dspark::graph::graph(const llama_model & model, const llm_graph_params & params) - : llama_model_dflash::graph(model, params) {} - -// anchor (committed last) token of every draft block: token 0 of each block in the ubatch -class llm_graph_input_dspark_anchor : public llm_graph_input_i { -public: - llm_graph_input_dspark_anchor(uint32_t block_size) : block_size(block_size) {} - virtual ~llm_graph_input_dspark_anchor() = default; - - void set_input(const llama_ubatch * ubatch) override { - GGML_ASSERT(ubatch->token); - const int64_t n_blocks = anchors->ne[0]; - std::vector buf(n_blocks); - for (int64_t j = 0; j < n_blocks; ++j) { - buf[j] = ubatch->token[j*block_size]; - } - ggml_backend_tensor_set(anchors, buf.data(), 0, n_blocks*sizeof(int32_t)); - } - - bool can_reuse(const llm_graph_params & params) override { - return params.ubatch.token && anchors && - anchors->ne[0]*(int64_t) block_size == (int64_t) params.ubatch.n_tokens; - } - - ggml_tensor * anchors = nullptr; // I32 [n_blocks] - - const uint32_t block_size; -}; - -// DSpark decoder: DFlash decoder + Markov bias on the draft logits, 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)) -template <> -llama_model_dspark::graph::graph(const llama_model & model, const llm_graph_params & params) - : llama_model_dflash::graph(model, params) { - // KV-injection (embd) batch: no logits to bias - if (ubatch.embd) { - return; - } - - ggml_tensor * w1 = model.dspark_markov_w1; - ggml_tensor * w2 = model.dspark_markov_w2; - GGML_ASSERT(w1 && w2 && "DSpark markov 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 int64_t bs = model.hparams.n_dspark_block; - GGML_ASSERT(bs > 0); - - // the drafting loop always submits whole anchor-first blocks - if (n_tok % bs != 0) { - return; - } - const int64_t n_blocks = n_tok / bs; - - auto inp = std::make_unique((uint32_t) bs); - inp->anchors = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_blocks); - ggml_set_input(inp->anchors); - ggml_tensor * prev = inp->anchors; // I32 [n_blocks] - res->add_input(std::move(inp)); - - ggml_tensor * cat = 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 < bs; ++i) { - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, ggml_get_rows(ctx0, 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, bs*base->nb[1], i*base->nb[1]); - ggml_tensor * col = ggml_add(ctx0, base_i, bias); - - cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; - - if (i + 1 < bs) { - prev = ggml_argmax(ctx0, col); // I32 [n_blocks] - } - } - - // cat is position-major; restore the ubatch's block-major order - ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, bs); - out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs, n_blocks] - out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); - - res->t_logits = out; - ggml_build_forward_expand(gf, out); -} diff --git a/src/models/models.h b/src/models/models.h index 034e1282e338..a86ae05aae6b 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1254,22 +1254,6 @@ struct llama_model_dflash : public llama_model_base { }; -struct llama_model_dspark : public llama_model_dflash { - llama_model_dspark(const struct llama_model_params & params) : llama_model_dflash(params) {} - // extend the DFlash hparams/tensors with the block size and the Markov / confidence heads - void load_arch_hparams(llama_model_loader & ml) override; - void load_arch_tensors(llama_model_loader & ml) override; - - // the DFlash graphs plus the in-graph Markov head on the decoder's draft logits - template - struct graph : public llama_model_dflash::graph { - graph(const llama_model & model, const llm_graph_params & params); - }; - - std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; -}; - - struct llama_model_mistral4 : public llama_model_deepseek2 { llama_model_mistral4(const struct llama_model_params & params) : llama_model_deepseek2(params) {} // reuse load_arch_hparams and load_arch_tensors from llama_model_deepseek2 diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index fb71e9b0ec9c..57d33a627d5d 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -457,7 +457,7 @@ static int save_models(const llm_arch target_arch, const size_t seed, const ggml if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } - if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH || arch == LLM_ARCH_DSPARK) { + if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { continue; } for (bool moe : {false, true}) { @@ -563,7 +563,7 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg if (arch == LLM_ARCH_GEMMA4 || arch == LLM_ARCH_GEMMA4_ASSISTANT) { continue; // FIXME: ISWA KV cache initialization needs more fixture params } - if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH || arch == LLM_ARCH_DSPARK) { + if (arch == LLM_ARCH_EAGLE3 || arch == LLM_ARCH_DFLASH) { continue; } From 5fbc2c2e5d74f251a801f1c534d6f486e56d77c3 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Thu, 9 Jul 2026 21:15:09 +0800 Subject: [PATCH 07/12] 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). --- common/arg.cpp | 8 ++++++++ common/common.h | 2 ++ common/speculative.cpp | 12 +++++++----- docs/speculative.md | 3 +++ src/models/dflash.cpp | 34 ++++++++++++++++++++++++++++++++-- tools/cli/README.md | 3 ++- tools/server/README.md | 3 ++- tools/server/server-schema.cpp | 4 ++++ 8 files changed, 60 insertions(+), 9 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 8a74e5f2b1f9..36b5b6c7e2a9 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -3866,6 +3866,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.p_min = std::stof(value); } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_P_MIN")); + add_opt(common_arg( + {"--spec-draft-conf-min"}, "P", + string_format("DSpark: minimum predicted acceptance from the draft confidence head to keep a " + "drafted token; truncates the block at the first position below it (0.0 = disabled) (default: %.2f)", (double)params.speculative.draft.conf_min), + [](common_params & params, const std::string & value) { + params.speculative.draft.conf_min = std::stof(value); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_CONF_MIN")); add_opt(common_arg( {"--spec-draft-backend-sampling"}, {"--no-spec-draft-backend-sampling"}, diff --git a/common/common.h b/common/common.h index ca02c2964198..29eb2d1074c2 100644 --- a/common/common.h +++ b/common/common.h @@ -331,6 +331,8 @@ struct common_params_speculative_draft { float p_split = 0.1f; // speculative decoding split probability float p_min = 0.0f; // minimum speculative decoding probability (greedy) + float conf_min = 0.0f; // DSpark: min predicted acceptance from the confidence head (0 = disabled) + bool backend_sampling = true; // offload draft sampling to the backend (default: on) common_params_model mparams; diff --git a/common/speculative.cpp b/common/speculative.cpp index 2b9836ba188a..d28e8e98aa6d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1215,7 +1215,7 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl_dra params.n_max = std::min(params_in.draft.n_max, block_size); LOG_INF("%s: adding speculative implementation 'draft-dspark'\n", __func__); - LOG_INF("%s: - block_size=%d, n_max=%d\n", __func__, block_size, params.n_max); + LOG_INF("%s: - block_size=%d, n_max=%d, conf_min=%.2f\n", __func__, block_size, params.n_max, params.conf_min); } void draft(common_speculative_draft_params_vec & dparams) override { @@ -1274,11 +1274,13 @@ struct common_speculative_impl_draft_dspark : public common_speculative_impl_dra const int32_t nb = n_block[seq_id]; // drafts to keep (<= block_size) auto * smpl = smpls[seq_id].get(); - // greedily read the predicted block at this sequence's noise positions 1..nb-1 - // TODO: implement confidence-scheduled prefix pruning (the "scheduled" part of DSpark): - // use the confidence head to truncate the drafted block at the first low-confidence - // position instead of always keeping all n_draft tokens + // greedily read the predicted block at this sequence's noise positions 1..nb-1. + const float * conf = params.conf_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; for (int32_t i = 0; i < nb; ++i) { + if (conf && conf[(beg + i) * n_embd_dec] < params.conf_min) { + break; + } + common_sampler_sample(smpl, ctx_dft, beg + i, true); const auto * cur_p = common_sampler_get_candidates(smpl, true); diff --git a/docs/speculative.md b/docs/speculative.md index 84d45d47c54d..3957db85c9c1 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -100,6 +100,9 @@ llama-server -m Qwen3-4B.gguf -md Qwen3-4B-DSpark.gguf \ `--spec-draft-n-max` is clamped to the draft model's trained block size. +`--spec-draft-conf-min P` truncates each drafted block at the first position whose predicted +acceptance (from the draft's confidence head, if present) falls below `P` (default 0 = disabled). + Currently only drafts with a Qwen3 backbone are supported; support for other backbones (e.g. Gemma4) is planned. diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 7894e9dfa49a..28fbc648e599 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -156,11 +156,17 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, bs*tokens->nb[0], 0); prev = ggml_cont_1d(ctx0, prev, n_blocks); // I32 [n_blocks] - ggml_tensor * cat = nullptr; + // confidence head (optional): predicted acceptance per position, from the same + // hidden state that feeds the lm_head plus the markov embedding of prev + ggml_tensor * h = model.dspark_conf_proj ? res->t_embd : nullptr; // [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 < bs; ++i) { - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, ggml_get_rows(ctx0, w1, prev)); // [n_vocab, n_blocks] + 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, bs*base->nb[1], i*base->nb[1]); @@ -168,6 +174,19 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; + if (h) { + // conf(i) = sigmoid(conf_proj . [h(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] + ggml_tensor * h_i = ggml_view_2d(ctx0, h, h->ne[0], n_blocks, bs*h->nb[1], i*h->nb[1]); + ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, h_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 < bs) { prev = ggml_argmax(ctx0, col); // I32 [n_blocks] } @@ -178,6 +197,17 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs, n_blocks] out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); + if (cat_conf) { + // same position-major -> block-major reorder as the logits + ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, bs); + conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); + conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); + + 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); } diff --git a/tools/cli/README.md b/tools/cli/README.md index 89be3a5b365e..d0555685784b 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -198,11 +198,12 @@ | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | +| `--spec-draft-conf-min P` | DSpark: minimum predicted acceptance from the draft confidence head to keep a drafted token(default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_CONF_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | -| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | +| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | | `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) | | `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) | | `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) | diff --git a/tools/server/README.md b/tools/server/README.md index 1941dcea202d..2419f090f484 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -254,11 +254,12 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | +| `--spec-draft-conf-min P` | DSpark: minimum predicted acceptance from the draft confidence head to keep a drafted token(default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_CONF_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | -| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | +| `--spec-type none,draft-simple,draft-eagle3,draft-mtp,draft-dflash,draft-dspark,ngram-simple,ngram-map-k,ngram-map-k4v,ngram-mod,ngram-cache` | comma-separated list of types of speculative decoding to use (default: none)

(env: LLAMA_ARG_SPEC_TYPE) | | `--spec-ngram-mod-n-min N` | minimum number of ngram tokens to use for ngram-based speculative decoding (default: 48) | | `--spec-ngram-mod-n-max N` | maximum number of ngram tokens to use for ngram-based speculative decoding (default: 64) | | `--spec-ngram-mod-n-match N` | ngram-mod lookup length (default: 24) | diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 89026eb4e3f0..13acd5cf73a2 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -209,6 +209,10 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_hard_limits(0.0f, 1.0f) ->set_desc("Minimum speculative decoding probability for draft tokens (0 = greedy)")); + add((new field_num("speculative.conf_min", params.speculative.draft.conf_min)) + ->set_hard_limits(0.0f, 1.0f) + ->set_desc("DSpark: minimum confidence-head acceptance to keep a drafted token (0 = disabled)")); + add((new field_str("speculative.type")) ->set_desc("Speculative decoding method (for debugging and research purposes)") ->set_handler([&](field_eval_context & ctx, const json & data) { From efd1ccfd1cb16eadd73bb89be96adc5eea3b819d Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Thu, 9 Jul 2026 23:28:57 +0800 Subject: [PATCH 08/12] fold the dspark impl into dflash, selected by spec type --- common/speculative.cpp | 149 ++++++++++------------------------------- src/models/dflash.cpp | 3 - 2 files changed, 37 insertions(+), 115 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index d28e8e98aa6d..dd2eb307e7ee 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -919,6 +919,9 @@ 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 shifted; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -929,6 +932,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) : common_speculative_impl(type, n_seq) , params(params.draft) + , shifted(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; @@ -955,16 +959,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: - 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: 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, conf_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min, this->params.conf_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, * (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, * (block_size-1)]: in-place denoising yields at most + // block_size-1 drafts, anchor-first (DSpark) blocks yield a full block_size + const int32_t n_draft_max = shifted ? 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); @@ -1133,10 +1139,11 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { n_draft = std::min(n_draft, dp.n_max); } - const int32_t n_block_tokens = n_draft + 1; // id_last + n_draft * + // DSpark blocks are always submitted whole (the Markov head keys anchors off block boundaries) + const int32_t n_submit = shifted ? block_size : n_draft + 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) { + n_block [seq_id] = n_draft; + for (int32_t i = 0; i < n_submit; ++i) { common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); } } @@ -1158,23 +1165,33 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } auto & dp = dparams[seq_id]; - const int32_t beg = i_block_beg[seq_id]; - const int32_t n_block_tokens = n_block[seq_id]; + const int32_t beg = i_block_beg[seq_id]; + const int32_t n_draft = n_block[seq_id]; auto * smpl = smpls[seq_id].get(); 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); + // DSpark drafts expose a per-position predicted acceptance through the nextn + // channel; --spec-draft-conf-min prunes the block at the first position below it + const float * conf = shifted && params.conf_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; + + // greedily read the predicted block: DSpark blocks predict the NEXT token from position 0 on + for (int32_t k = 0; k < n_draft; ++k) { + const int32_t i = beg + k + (shifted ? 0 : 1); + + if (conf && conf[(size_t) i * n_embd_dec] < params.conf_min) { + break; + } + + common_sampler_sample(smpl, ctx_dft, i, true); const auto * cur_p = common_sampler_get_candidates(smpl, true); - for (int k = 0; k < std::min(3, (int) cur_p->size); ++k) { + for (int j = 0; j < std::min(3, (int) cur_p->size); ++j) { 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()); + seq_id, j, k, cur_p->data[j].id, cur_p->data[j].p, + common_token_to_piece(ctx_dft, cur_p->data[j].id).c_str()); } const llama_token id = cur_p->data[0].id; @@ -1203,98 +1220,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } }; -// DSpark: DFlash backbone + a semi-autoregressive Markov head; reuses process(), only draft() differs. -// The draft model is a DFlash GGUF carrying the Markov head tensors -- the decoder graph picks them -// up by presence, this impl only changes how the noise blocks are submitted and read back. -struct common_speculative_impl_draft_dspark : public common_speculative_impl_draft_dflash { - common_speculative_impl_draft_dspark(const common_params_speculative & params_in, uint32_t n_seq) - : common_speculative_impl_draft_dflash(params_in, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) - { - // DSpark drafts from position 0 of the block (the anchor position), so unlike DFlash - // it can draft a full block_size tokens per step; undo the base class's n_max clamp - params.n_max = std::min(params_in.draft.n_max, block_size); - - LOG_INF("%s: adding speculative implementation 'draft-dspark'\n", __func__); - LOG_INF("%s: - block_size=%d, n_max=%d, conf_min=%.2f\n", __func__, block_size, params.n_max, params.conf_min); - } - - void draft(common_speculative_draft_params_vec & dparams) override { - auto * ctx_dft = params.ctx_dft; - - common_batch_clear(batch); - - std::vector i_block_beg(n_seq, -1); - std::vector n_block (n_seq, 0); - - for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { - auto & dp = dparams[seq_id]; - if (!dp.drafting) { - continue; - } - - common_sampler_reset(smpls[seq_id].get()); - - 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); - } - n_draft = std::min(n_draft, block_size); - if (n_draft <= 0) { - continue; - } - - // anchor-first block [id_last, * (block_size-1)]: submit the whole block so the - // in-graph Markov head can key anchors off the block boundaries; keep the first n_draft - i_block_beg[seq_id] = batch.n_tokens; - n_block [seq_id] = n_draft; - for (int32_t i = 0; i < block_size; ++i) { - common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); - } - } - - if (batch.n_tokens == 0) { - return; - } - - if (llama_decode(ctx_dft, batch) != 0) { - LOG_WRN("%s: llama_decode failed\n", __func__); - return; - } - - for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { - if (i_block_beg[seq_id] < 0) { - continue; - } - auto & dp = dparams[seq_id]; - auto & result = *dp.result; - - const int32_t beg = i_block_beg[seq_id]; - const int32_t nb = n_block[seq_id]; // drafts to keep (<= block_size) - - auto * smpl = smpls[seq_id].get(); - // greedily read the predicted block at this sequence's noise positions 1..nb-1. - const float * conf = params.conf_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; - for (int32_t i = 0; i < nb; ++i) { - if (conf && conf[(beg + i) * n_embd_dec] < params.conf_min) { - break; - } - - common_sampler_sample(smpl, ctx_dft, beg + i, true); - - const auto * cur_p = common_sampler_get_candidates(smpl, true); - - const llama_token id = cur_p->data[0].id; - - common_sampler_accept(smpl, id, true); - - result.push_back(id); - } - } - } -}; - struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_params_speculative_draft params; // reuses the draft-model params slot (ctx_tgt/ctx_dft) @@ -2510,7 +2435,7 @@ common_speculative * common_speculative_init(common_params_speculative & params, break; } case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: { - impls.push_back(std::make_unique(config.params, n_seq)); + impls.push_back(std::make_unique(config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)); break; } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 28fbc648e599..dbdabede0ead 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -47,9 +47,6 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { R, n_vocab }, 0); dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { R, n_vocab }, 0); - // TODO: the confidence head is loaded but not yet used by the graph -- it should gate - // the draft length per block (early-exit on low confidence) instead of always - // drafting the full block dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + R, 1 }, TENSOR_NOT_REQUIRED); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); From acf49d98cba4abfe875f6bf4786bfc66766dc636 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Fri, 10 Jul 2026 11:54:03 +0800 Subject: [PATCH 09/12] address review comments --- common/speculative.cpp | 83 +++++++++++++++++++++++++++--------------- src/models/dflash.cpp | 29 +++++++++------ 2 files changed, 70 insertions(+), 42 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index dd2eb307e7ee..1f530eeb9bc5 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -920,7 +920,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { llama_token mask_token_id = 0; // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout - const bool shifted; + const bool is_dspark; const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -932,7 +932,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) : common_speculative_impl(type, n_seq) , params(params.draft) - , shifted(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) + , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; @@ -965,7 +965,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // DFlash input is [id_last, * (block_size-1)]: in-place denoising yields at most // block_size-1 drafts, anchor-first (DSpark) blocks yield a full block_size - const int32_t n_draft_max = shifted ? block_size : block_size - 1; + 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); @@ -1139,11 +1139,10 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { n_draft = std::min(n_draft, dp.n_max); } - // DSpark blocks are always submitted whole (the Markov head keys anchors off block boundaries) - const int32_t n_submit = shifted ? block_size : n_draft + 1; + 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_draft; - for (int32_t i = 0; i < n_submit; ++i) { + 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); } } @@ -1165,44 +1164,68 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } auto & dp = dparams[seq_id]; - const int32_t beg = i_block_beg[seq_id]; - const int32_t n_draft = n_block[seq_id]; + const int32_t beg = i_block_beg[seq_id]; + const int32_t n_block_tokens = n_block[seq_id]; auto * smpl = smpls[seq_id].get(); auto & result = *dp.result; - // DSpark drafts expose a per-position predicted acceptance through the nextn - // channel; --spec-draft-conf-min prunes the block at the first position below it - const float * conf = shifted && params.conf_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; + 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.conf_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; - // greedily read the predicted block: DSpark blocks predict the NEXT token from position 0 on - for (int32_t k = 0; k < n_draft; ++k) { - const int32_t i = beg + k + (shifted ? 0 : 1); + for (int32_t i = 0; i < n_block_tokens; ++i) { + const int32_t idx = beg + i; - if (conf && conf[(size_t) i * n_embd_dec] < params.conf_min) { - break; - } + if (conf && conf[(size_t) idx * n_embd_dec] < params.conf_min) { + break; + } - common_sampler_sample(smpl, ctx_dft, i, true); + common_sampler_sample(smpl, ctx_dft, idx, true); - const auto * cur_p = common_sampler_get_candidates(smpl, true); + const auto * cur_p = common_sampler_get_candidates(smpl, true); - for (int j = 0; j < std::min(3, (int) cur_p->size); ++j) { - LOG_DBG(" - seq_id %d, draft candidate %3d, pos %3d: %6d (%8.3f) '%s'\n", - seq_id, j, k, cur_p->data[j].id, cur_p->data[j].p, - common_token_to_piece(ctx_dft, cur_p->data[j].id).c_str()); - } + 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; + const llama_token id = cur_p->data[0].id; - if (cur_p->data[0].p < params.p_min) { - break; + if (cur_p->data[0].p < params.p_min) { + break; + } + + 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); - common_sampler_accept(smpl, id, true); + const auto * cur_p = common_sampler_get_candidates(smpl, true); - result.push_back(id); + 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); + } } if (result.size() < (size_t) params.n_min) { diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index dbdabede0ead..4991bb7e21b2 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -137,20 +137,25 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & const int64_t n_vocab = base->ne[0]; const int64_t n_tok = base->ne[1]; - // the trained draft block size, in tokens (anchor + n-1 masks) + // the trained draft block size bounds the runtime block length (anchor + up to bs-1 masks) 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 bs = std::stoi(it->second); GGML_ASSERT(bs > 0); - // the drafting loop always submits whole anchor-first blocks - if (n_tok % bs != 0) { + // the drafting loop submits one uniform anchor-first block per sequence; recover the runtime + // block length from the ubatch instead of assuming the full trained block_size + const int64_t n_blocks = g.ubatch.n_seqs_unq; + if (n_blocks == 0 || n_tok % n_blocks != 0) { + return; + } + const int64_t bs_rt = n_tok / n_blocks; // runtime block length, <= bs + if (bs_rt > bs) { return; } - const int64_t n_blocks = n_tok / bs; // anchor (committed last) token of every block: token 0 of each block, i.e. a strided view - ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, bs*tokens->nb[0], 0); + ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, bs_rt*tokens->nb[0], 0); prev = ggml_cont_1d(ctx0, prev, n_blocks); // I32 [n_blocks] // confidence head (optional): predicted acceptance per position, from the same @@ -161,19 +166,19 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & 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 < bs; ++i) { + for (int64_t i = 0; i < bs_rt; ++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, bs*base->nb[1], i*base->nb[1]); + ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, bs_rt*base->nb[1], i*base->nb[1]); ggml_tensor * col = ggml_add(ctx0, base_i, bias); cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; if (h) { // conf(i) = sigmoid(conf_proj . [h(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] - ggml_tensor * h_i = ggml_view_2d(ctx0, h, h->ne[0], n_blocks, bs*h->nb[1], i*h->nb[1]); + ggml_tensor * h_i = ggml_view_2d(ctx0, h, h->ne[0], n_blocks, bs_rt*h->nb[1], i*h->nb[1]); ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, h_i), w1_prev, 0); ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); if (model.dspark_conf_proj_b) { @@ -184,19 +189,19 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; } - if (i + 1 < bs) { + if (i + 1 < bs_rt) { prev = ggml_argmax(ctx0, col); // I32 [n_blocks] } } // cat is position-major; restore the ubatch's block-major order - ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, bs); - out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs, n_blocks] + ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, bs_rt); + out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs_rt, n_blocks] out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); if (cat_conf) { // same position-major -> block-major reorder as the logits - ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, bs); + ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, bs_rt); conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); From c946653aa9218ad90fa0126668c38a565c429f87 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 10 Jul 2026 15:36:10 +0000 Subject: [PATCH 10/12] dspark: clean up and improve naming --- common/speculative.cpp | 9 ++--- gguf-py/gguf/constants.py | 9 ++--- src/models/dflash.cpp | 70 +++++++++++++++++++-------------------- 3 files changed, 42 insertions(+), 46 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 1f530eeb9bc5..e2480cac5cdb 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -964,7 +964,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { 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, * (block_size-1)]: in-place denoising yields at most - // block_size-1 drafts, anchor-first (DSpark) blocks yield a full block_size + // 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", @@ -1195,10 +1195,6 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { 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); @@ -2458,7 +2454,8 @@ common_speculative * common_speculative_init(common_params_speculative & params, break; } case COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK: { - impls.push_back(std::make_unique(config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)); + impls.push_back(std::make_unique( + config.params, n_seq, COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK)); break; } case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c1895c0fe0b8..d93c87f110da 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -954,9 +954,10 @@ class MODEL_TENSOR(IntEnum): # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping - DSPARK_MARKOV_W1 = auto() # dspark markov head: prev-token embed - DSPARK_MARKOV_W2 = auto() # dspark markov head: bias projection - DSPARK_CONF_PROJ = auto() # dspark confidence head: proj + # dspark + DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed + DSPARK_MARKOV_W2 = auto() # markov head: bias projection + DSPARK_CONF_PROJ = auto() # confidence head # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -4242,7 +4243,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, - # optional DSpark heads (present only in DSpark drafts) + # optional DSpark heads MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 4991bb7e21b2..7c716f3c76b7 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -36,21 +36,21 @@ 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; detect it by tensor presence + // 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 R = markov_meta->ne[0]; // markov rank + const int64_t dspark_markov_rank = markov_meta->ne[0]; - dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { R, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { R, n_vocab }, 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 + R, 1 }, TENSOR_NOT_REQUIRED); + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED); 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) R); + 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); @@ -121,10 +121,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_grap ggml_build_forward_expand(gf, cur); } -// DSpark drafts only (DFlash + Markov head): Markov bias on the draft logits, -// 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)) +// 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; @@ -137,49 +134,51 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & const int64_t n_vocab = base->ne[0]; const int64_t n_tok = base->ne[1]; - // the trained draft block size bounds the runtime block length (anchor + up to bs-1 masks) 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 bs = std::stoi(it->second); - GGML_ASSERT(bs > 0); + const int64_t block_size = std::stoi(it->second); + GGML_ASSERT(block_size > 0); - // the drafting loop submits one uniform anchor-first block per sequence; recover the runtime - // block length from the ubatch instead of assuming the full trained block_size const int64_t n_blocks = g.ubatch.n_seqs_unq; if (n_blocks == 0 || n_tok % n_blocks != 0) { return; } - const int64_t bs_rt = n_tok / n_blocks; // runtime block length, <= bs - if (bs_rt > bs) { + // 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 - ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, bs_rt*tokens->nb[0], 0); - prev = ggml_cont_1d(ctx0, prev, n_blocks); // I32 [n_blocks] + const size_t token_stride = (size_t) block_drafts * tokens->nb[0]; + const size_t base_stride = (size_t) block_drafts * base->nb[1]; - // confidence head (optional): predicted acceptance per position, from the same - // hidden state that feeds the lm_head plus the markov embedding of prev - ggml_tensor * h = model.dspark_conf_proj ? res->t_embd : nullptr; // [n_embd, n_tok] + ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); + prev = ggml_cont_1d(ctx0, prev, n_blocks); + + // optional confidence head: predicts per-position acceptance + ggml_tensor * conf_inp = model.dspark_conf_proj ? res->t_embd : nullptr; // [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 < bs_rt; ++i) { + 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, bs_rt*base->nb[1], i*base->nb[1]); + 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; - if (h) { - // conf(i) = sigmoid(conf_proj . [h(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] - ggml_tensor * h_i = ggml_view_2d(ctx0, h, h->ne[0], n_blocks, bs_rt*h->nb[1], i*h->nb[1]); - ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, h_i), w1_prev, 0); + if (conf_inp) { + // 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); @@ -189,19 +188,18 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; } - if (i + 1 < bs_rt) { - prev = ggml_argmax(ctx0, col); // I32 [n_blocks] + if (i + 1 < block_drafts) { + prev = ggml_argmax(ctx0, col); } } - // cat is position-major; restore the ubatch's block-major order - ggml_tensor * out = ggml_reshape_3d(ctx0, cat, n_vocab, n_blocks, bs_rt); - out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, bs_rt, n_blocks] + // 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); if (cat_conf) { - // same position-major -> block-major reorder as the logits - ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, bs_rt); + 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); @@ -386,7 +384,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_build_forward_expand(gf, cur); - // DSpark drafts only: bias the draft logits with the Markov head + // DSpark: bias the draft logits with the Markov head if (model.dspark_markov_w1) { build_dspark_markov_head(*this, model, inp_tokens); } From be954b2b5a1fbac2c451bdc15cbd9cf6ae197273 Mon Sep 17 00:00:00 2001 From: wjinxu <1299461899@qq.com> Date: Fri, 17 Jul 2026 21:10:44 +0800 Subject: [PATCH 11/12] update readme --- tools/cli/README.md | 2 +- tools/server/README.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/cli/README.md b/tools/cli/README.md index d0555685784b..708fa9c19cde 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -198,7 +198,7 @@ | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | -| `--spec-draft-conf-min P` | DSpark: minimum predicted acceptance from the draft confidence head to keep a drafted token(default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_CONF_MIN) | +| `--spec-draft-conf-min P` | DSpark: minimum predicted acceptance from the draft confidence head to keep a drafted token; truncates the block at the first position below it (0.0 = disabled) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_CONF_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | diff --git a/tools/server/README.md b/tools/server/README.md index 2419f090f484..5bda4e5ebf60 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -254,7 +254,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | -| `--spec-draft-conf-min P` | DSpark: minimum predicted acceptance from the draft confidence head to keep a drafted token(default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_CONF_MIN) | +| `--spec-draft-conf-min P` | DSpark: minimum predicted acceptance from the draft confidence head to keep a drafted token; truncates the block at the first position below it (0.0 = disabled) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_CONF_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | From aa3a4feab0966316037d58b0ace5139f8b24fe32 Mon Sep 17 00:00:00 2001 From: Ruixiang Wang Date: Fri, 10 Jul 2026 15:58:56 +0000 Subject: [PATCH 12/12] remove trailing whitespace --- src/models/dflash.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 7c716f3c76b7..6afb6a8523f4 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -161,7 +161,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & 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) {