spec: add DSpark speculative decoding#25173
Conversation
This comment was marked as resolved.
This comment was marked as resolved.
wjinxu
requested review from
a team,
CISC,
JohannesGaessler and
ggerganov
as code owners
|
Hi @CISC @ggerganov , this adds DSpark speculative decoding on top of the merged DFlash drafter (#22105). It's a small change — a new dspark draft arch and draft-dspark spec type that reuse DFlash's graph, feature extraction, KV-cache injection and verify path unchanged; the only new logic is the semi-autoregressive Markov head in draft(). Greedy stays lossless. I benchmarked it against the merged DFlash using DeepSeek's released Qwen3 DSpark drafts. On Qwen3-8B at bf16 / Q8_0 / Q4_K_M, DSpark beats DFlash on every domain (e.g. GSM8K bf16 4.06× vs 3.12×; full per-domain tables in the PR description). I believe it's ready for review and I'm happy to walk through any part of it. |
|
Can you run SpeedBench to do the full comparison between DFlash and DSpark with the same |
|
@ruixiang63 I've run the SpeedBench test set as you suggested, and updated the results in the PR description. DSpark does outperform DFlash across the board. |
wjinxu
force-pushed
the
dspark-upstream
branch
2 times, most recently
from
aae2941 to
d8b38f2
Compare
|
could you give some examples how to use? |
Good point — I've updated the PR description with a more detailed, copy-pasteable end-to-end example (download → convert → build → run → curl). Let me know if anything's unclear. |
|
DSV4 support was merged in #24162, ideally this PR should cover that model as well and try to replicate a similar speedup |
Thanks! DeepSeek hasn't open-sourced the DSpark weights for DeepSeek-V4 though — only the Qwen3 and Gemma4 drafts are released. So this PR covers Qwen3 for now, and I'll add Gemma4 as a small follow-up. |
|
I think they're a part of the spec decoding module https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash-DSpark i.e not distributed separately |
Sorry, and thanks for the heads-up. For this PR I'd like to keep the scope a bit narrower for now - land the Qwen3 DSpark path first and get it solid, then add Gemma4 and DSV4 as follow-ups. Does that sound ok? |
|
Okay, will try to review. From a cursory look it does not look like adding |
I agree that |
|
@am17an Thanks for the feedback — you were right about the API. I've reworked the PR: All llama_dspark_* public APIs are gone. The Markov head is now applied inside the dspark decoder graph(chained per block position, vectorized across blocks), so llama.h / llama-ext.h are untouched and common/speculative.cpp only overrides draft() on top of the DFlash impl, same pattern as before. I've reviewed and can explain every line of the code myself, and re-verified. I'd really appreciate another look when you have time. |
github-actions
Bot
added
the
documentation
|
@ruixiang63 Could you take another look at this PR when you have time? Thanks! |
|
I spent this morning testing wjinxu's DSpark speculative decoding implementation (many thanks for his hard work) and I can confirm that it indeed works as wjinxu described. See the attached writeup. A few remarks:
@wjinxu : please, could you add the code for Gemma4 support? I would be more than willing to test it on my Linux workstation. Thank you! Attached file is a long description of the testing session: |
am17an
left a comment
am17an
left a comment
There was a problem hiding this comment.
Sorry I forgot about this PR. It looks okay, but it should have some TODOs in code about how incomplete this implementation is. From what I understand it just adds the autoregressive layer over DFlash which helps in the keeping the generated tokens conditioned on the previous generated tokens. cc @ruixiang63 @ggerganov
Thanks for the thorough testing and the detailed writeup — much appreciated! Glad to hear DSpark works as described on your setup. |
|
Thanks for taking a look! You're right — the current implementation only adds the autoregressive layer on top of DFlash, so the drafted tokens stay conditioned on previously generated ones, but the rest is still incomplete. I'll add TODO comments in the code to make the missing pieces explicit. |
- 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)
I've added TODOs marking the incomplete parts, see 1829020. The list:
I plan to tackle backbone support (Gemma4) first: with the autoregressive Markov head, DSpark already outperforms DFlash at essentially the same draft cost, so the current state is a net win as-is. The confidence-scheduled pruning would come after that. |
@am17an Which parts are missing - I haven't looked at DSpark yet? |
|
DSpark also changes the number of verification tokens based on another thing called the "Confidence Head". From the image you can see the verifier drops verifying the last token (C4) based on some confidence thresholds + server load. 
|
|
Thanks for raising the PR, and sorry for the delay in reviewing it. I have a few suggestions:
Thanks for your work on this. Let’s explore how we can improve it further and keep things moving forward. |
|
|
||
| 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 |
There was a problem hiding this comment.
Judging from the comment, DFlash already has block size, why we need to extend with block_size here?
There was a problem hiding this comment.
Removed — DSpark no longer has its own block size handling.
| @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)) |
There was a problem hiding this comment.
This convert code is 95% similar to DFlash. Can we merge it into DFlash to minimize the scope of changes?
| : 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) |
There was a problem hiding this comment.
why do we need to change here?
There was a problem hiding this comment.
Because the dspark impl inherits from the dflash impl, and type is a const member — it can only be set via the constructor, so it has to be passed through here.
| // 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) { |
There was a problem hiding this comment.
I guess this is the only different part from DFlash in current implementation, right?
| 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<int32_t> 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; | ||
| }; |
There was a problem hiding this comment.
Is this logic should be same as DFlash? I feel it is not the correct way to do it here.
There was a problem hiding this comment.
You're right, Reworked it to take a strided view of that input instead.
|
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.
I see. Then let's focus on Qwen3 now.
I think you’ve already loaded I checked the new updates, and it looks better. Thanks for the quick turnaround! I’m still not sure whether we need a separate DSpark speculative decoding class, or whether we should integrate it directly into the current DFlash speculative decoding implementation, conditioned on the presence of the Markov head and confidence head. |
Agreed. I'll add a static flag for this.
While looking into this I found that the presence of the Markov tensors cannot be used to tell DSpark and DFlash drafts apart. DeepSeek's released DFlash checkpoints (e.g. deepseek-ai/dflash_qwen3_8b_block7) are trained with the DSpark architecture (Qwen3DSparkModel) but with markov_rank = 0, so they carry no Markov tensors — yet they still use DSpark's anchor-first block layout (drafts are read from position 0 of the block, a full block_size per step), which differs from z-lab's DFlash drafts (read from position 1, block_size - 1 per step). |
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).
I understand, but this is DFlash rather than DSpark, right? I think the scope of this PR is to add DSpark support. I don’t see much value in adding DeepSeek’s DFlash model at this moment, especially since the DFlash models are based on the outdated Qwen3 models. |
8 participants
This PR adds DSpark speculative decoding, layered on the merged DFlash drafter. DSpark (DeepSeek + PKU, 2026 — "Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation", the DeepSpec repo) is DFlash plus a small semi-autoregressive Markov head: where DFlash takes an independent argmax at each block position (every position marginalizes over all possible predecessors, so acceptance decays along the block), DSpark adds a low-rank, previous-token-conditioned logit bias and samples the block left-to-right, so each draft conditions on the one actually sampled before it. This lifts accepted length at near-zero extra draft cost.
DSpark reuses the entire DFlash machinery unchanged — the encoder/decoder graph, target-layer feature extraction (
llama_set_embeddings_layer_inp/_nextn), KV-cache injection, and the verify/accept path. The only additions are:dspark(llama_model_dspark : llama_model_dflash) that reuses the DFlash graph and additionally loads the Markov head (markov_w1,markov_w2) and an optional confidence head; it shares the target's token-embeddings / lm_head (same as DFlash);draft-dspark(common_speculative_impl_draft_dspark : common_speculative_impl_draft_dflash) that reusesprocess()(extraction + injection) and overrides onlydraft(): the block is anchor-first (position 0 already predicts the first draft token) and sampled with the Markov biasbias(prev) = markov_w2 · markov_w1[prev], computed on-device (llama_dspark_markov_bias);Qwen3DSparkModelconverter.Greedy decoding is lossless: the Markov bias only changes which tokens are proposed; every draft is still verified against the target, so the output is identical to non-speculative greedy.
The confidence head is converted/loaded but not used at inference in this PR (phase 1); the draft-quality win from the Markov head is self-contained and is what the numbers below measure.
How to run
Complete example from scratch (Qwen3-8B). Drafts for other sizes are on the same org:
deepseek-ai/dspark_qwen3_{4b,8b,14b}_block7.1. Get the models — target + its DSpark draft:
2. Convert to GGUF — the draft ships no tokenizer and reuses the target's, so pass
--target-model-dir:python convert_hf_to_gguf.py Qwen3-8B --outtype bf16 --outfile Qwen3-8B.gguf python convert_hf_to_gguf.py dspark_qwen3_8b --outtype bf16 \ --target-model-dir Qwen3-8B --outfile Qwen3-8B-DSpark.ggufYou may quantize the target (e.g.
llama-quantize Qwen3-8B.gguf Qwen3-8B-Q4_K_M.gguf Q4_K_M); keep the draft bf16 — it's tiny, and acceptance is unaffected by target quant.3. Build with CUDA:
cmake -B build -DGGML_CUDA=ON && cmake --build build --config Release -j4. Run — the only DSpark-specific flags are
-md <draft>and--spec-type draft-dspark(
--spec-draft-n-max= draft tokens per step; the released checkpoints use block size 7):./build/bin/llama-server -m Qwen3-8B.gguf -md Qwen3-8B-DSpark.gguf \ --spec-type draft-dspark --spec-draft-n-max 7 \ --temp 0 --top-k 1 -np 1 -c 4096 -ngl 99 -fa on --jinja5. Send a request (the server logs
draft acceptance = ...per request):llama-cliworks the same way (-m ... -md ... --spec-type draft-dspark). Note:draft-dsparkneeds the target's hidden states (KV-cache injection), so usellama-server/llama-cli— thespeculative-simpleexample does not drive that path.Performance
SpeedBench (llama.cpp's own
tools/server/bench/speed-bench)Qwen3-8B (bf16), matched
--spec-draft-n-max 7,qualitativesplit (11 categories), greedy. Baseline is the same server with no draft model. DSpark reaches 1.88× overall decode speedup vs baseline (DFlash is 1.55×), and beats the merged DFlash on every one of the 11 categories (overall 1.21×).DSpark vs baseline:
DSpark vs the merged DFlash (same
--spec-draft-n-max 7):Hardware: RTX 4090. Target Qwen/Qwen3-8B (bf16), draft deepseek-ai/dspark_qwen3_8b_block7 (bf16). Greedy (
--temp 0 --top-k 1), no-thinking,--spec-draft-n-max 7. Baseline = same llama-server with no draft model. DFlash is the merged drafter (z-lab/Qwen3-8B-DFlash, b16), run at the same matched draft size for an apples-to-apples comparison. Per-domain aggregate over the listed prompt counts.Losslessness
Greedy decoding is lossless by construction (the draft is verified against the target). Output is coherent and matches non-speculative greedy.
Qwen3-4B, target bf16
DSpark vs baseline (DFlash was not benchmarked at 4B — no nested-schema 4B DFlash draft available):
Qwen3-8B, target bf16
Qwen3-8B, target Q8_0
Qwen3-8B, target Q4_K_M
DSpark beats the merged DFlash on every domain (higher accept rate and higher throughput), for a ~1.16× geomean speedup over DFlash. The gains are largest on reasoning (GSM8K +25pp accept, 1.30× over DFlash) and open chat (MT-Bench, 1.29×); on code (HumanEval) the two are close as both already accept ~80%.
Future work
Requirements