dspark: support speculators-format checkpoints - #26275
Conversation
github-actions
Bot
added
documentation
083c33f to
77e9a31
Compare
Yes, I’ve tested this checkpoint, but it didn’t perform very well. Its acceptance rate was roughly the same as dflash’s. |
77e9a31 to
e107846
Compare
| } | ||
| mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); | ||
|
|
||
| if (is_dspark) { |
There was a problem hiding this comment.
same here. Feel it is not necessary. Would be better to keep the code change minimal and simple.
|
What is the difference between this PR and #25549? Which one should we review first? |
|
I completely understand what you mean. How about we handle this one first? Since this isn't about adding a new backbone but adapting to another training framework, I think it should be relatively simpler. |
|
As for how to adapt the backbone going forward — Qwen 3.8 may be released in a few days — I think waiting to see how the community chooses could be a good approach. If the community goes with a separate backbone, we can adjust accordingly in time. |
|
can you rebase this PR? @wjinxu Now I have some bandwidth to review the details more closely. |
Speculators-format DSpark drafts (e.g. SpecForge exports for the Gemma-4-26B-A4B target) differ from the dense DeepSpec checkpoints in three ways: - the config nests the backbone hparams under transformer_layer_config and gives the extract layers as aux_hidden_state_layer_ids - the block is the DFlash 1+N fill-in layout: the anchor slot is a bonus token, not a prediction slot. Written as dflash.bonus_anchor; such drafts build the block and read the mask positions exactly like DFlash (n_max drafts from a 1+n_max block), only the Markov/confidence sampling comes from DSpark - the draft output vocab may be reduced (draft_vocab_size < vocab_size) with a d2t remap table. The converter expands lm_head/markov_w2 back to the full vocab and synthesizes an lm_head bias of -1e9 on the rows the draft cannot produce, so the runtime needs no d2t remapping. Such drafts ship their own (now optional) token_embd/output tensors instead of sharing the target's Verified against gemma4-26b-a4b-dspark: greedy outputs are byte-identical with and without the draft; acceptance 0.46, mean draft len 3.7 (n_max 6). Co-authored-by: desovo7 <942845546@qq.com> Assisted-by: Claude Fable 5
One class now covers every DSpark variant. What used to pick the class is a single flag, because the arch name turns out to be the only thing that separates the two families: SpecForge also exports a flat schema that carries no speculators_* fields yet still uses the 1+N bonus-anchor block, so keying on those fields would silently mis-read its drafts. Also rename i0 to i_first_pred in the draft read loop and the Markov head, and give the head a real bonus_anchor bool instead of testing i0 > 0. Converting the Qwen3-8B DeepSpec draft and both gemma-4 speculators drafts produces byte-identical GGUFs. The one behaviour change is that the markov_head_type check now also covers the DeepSpec checkpoints, which previously skipped it. Co-authored-by: desovo7 <942845546@qq.com> Assisted-by: Claude Opus 5
f806441 to
72f87af
Compare
Done. Thanks very much. |
| const bool is_dspark; | ||
|
|
||
| // dspark speculators exports: DFlash 1+N block, the anchor is not a prediction slot | ||
| bool bonus_anchor = false; |
There was a problem hiding this comment.
| bool bonus_anchor = false; | |
| bool sample_from_anchor = false; |
| return TextModel.filter_tensors(item) | ||
| return super().filter_tensors(item) | ||
|
|
||
| def _expand_reduced_vocab(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]: |
There was a problem hiding this comment.
Can you check how Eagle3 uses d2t and make this consistent with the existing implementation?
| yield from super().modify_tensors(bias, "lm_head.bias", None) | ||
|
|
||
| def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: | ||
| if name == "d2t": |
| hparams = ModelBase.load_hparams(dir_model, False) | ||
|
|
||
| # only the arch name separates the SpecForge drafts from the DeepSpec ones | ||
| self._is_specforge = hparams["architectures"][0] in ("DSparkDraftModel", "DSparkSpeculator") |
There was a problem hiding this comment.
I don’t think we need _is_specforge here. The Eagle3 conversion function can handle checkpoints trained with SpecForge as well as many other trained Eagle3 variants. I’d prefer to avoid adding framework-specific training logic or checks here.
There was a problem hiding this comment.
Renamed _is_specforge to _sample_from_anchor, determined by which architecture is in use. However, I think deciding whether to sample based on the architecture is necessary behavior.
| name, gen = item | ||
| if name.endswith(("embed_tokens.weight", "lm_head.weight")): | ||
| name = item[0] | ||
| if name == "t2d": # training-only target->draft mask |
There was a problem hiding this comment.
as far as I know we don't even need to mention t2d here, right?
| TARGET_LAYERS = "{arch}.target_layers" | ||
| TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" | ||
| BLOCK_SIZE = "{arch}.block_size" | ||
| BONUS_ANCHOR = "{arch}.bonus_anchor" |
There was a problem hiding this comment.
| BONUS_ANCHOR = "{arch}.bonus_anchor" | |
| SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" |
| GGML_ASSERT(block_size > 0); | ||
|
|
||
| // bonus anchor (SpecForge exports): slot 0 holds a bonus token, not a prediction slot | ||
| const auto it_ba = model.gguf_kv.find("dflash.bonus_anchor"); |
There was a problem hiding this comment.
Please use the new name and improve it_ba naming here.
There was a problem hiding this comment.
Done. Renamed it_ba to it_sample_from_anchor.
- rename bonus_anchor to sample_from_anchor (GGUF key and code), matching the checkpoint config field; absent key still means anchor-first - rework the reduced draft vocab to match EAGLE3: d2t is written as I64 absolute target ids and the logits are scattered at runtime, instead of expanding lm_head/markov_w2 and synthesizing an output bias at conversion - move the t2d skip to modify_tensors, like EAGLE3 - drop _is_specforge: the arch name only picks the sample_from_anchor default, embed/lm_head sharing is decided by the draft vocab size - deduplicate the tok_embd create_tensor left behind by the rebase Verified with the RedHat gemma-4-31b speculator draft: greedy output is byte-identical with and without the draft; acceptance 0.26 (n_max 7). Co-authored-by: desovo7 <942845546@qq.com> Assisted-by: Claude Fable 5
2 participants
Overview
Follow-up to #25173. This adds support for DSpark drafts exported in the speculators format (SpecForge / RedHat), which vLLM gained in vllm-project/vllm#47093.
Format differences
These checkpoints differ from the dense DeepSpec checkpoints in same ways:
The speculators config uses
sample_from_anchorto declare the block layout:trueselects anchor-first, matching dense DeepSpec where every slot predicts a token;falseselects DFlash-style1+Ninfilling, where the anchor slot contains the bonus token and1 + n_maxslots produce n_max draft tokens. A missing field defaults to false, since legacy exports only use the1+Nlayout.A checkpoint may use a pruned draft output vocabulary
(draft_vocab_size < vocab_size)together with a draft-to-target(d2t)remapping table.Verification
makora-ai/gemma4-26b-a4b-dsparkRedHatAI/gemma-4-31B-it-speculator.dsparkAdditional information
Testing and feedback are welcome.
Requirements