diff --git a/common/speculative.cpp b/common/speculative.cpp index cd2dfc760b88..ae55e357d51a 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -926,6 +926,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout const bool is_dspark; + // dspark speculators + bool sample_from_anchor = true; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -960,16 +963,20 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { if (llama_model_meta_val_str(model_dft, "dflash.block_size", buf, sizeof(buf)) >= 0) { block_size = std::atoi(buf); } + if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) { + sample_from_anchor = std::strcmp(buf, "true") == 0; + } } 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: - 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); + LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, + block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false"); // DFlash input is [id_last, * (block_size-1)]: in-place denoising yields at most - // block_size-1 draft tokens, DSpark yield a full block_size draft tokens - const int32_t n_draft_max = is_dspark ? block_size : block_size - 1; + // block_size-1 draft tokens, anchor-first DSpark yields a full block_size draft tokens + const int32_t n_draft_max = is_dspark && sample_from_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); @@ -1175,7 +1182,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n_draft = params.n_max; - const int32_t n_block_tokens = n_draft + (is_dspark ? 0 : 1); + const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1); i_block_beg[seq_id] = batch.n_tokens; n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { @@ -1208,11 +1215,11 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; if (is_dspark) { - // DSpark predicts the next token from position 0 and optionally truncates - // at the first position below the confidence threshold. + // DSpark: read from the first draft slot, truncate below the confidence threshold const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; - - for (int32_t i = 0; i < n_block_tokens; ++i) { + // bonus-anchor drafts read the mask positions only, like DFlash + const int32_t i_draft_beg = sample_from_anchor ? 0 : 1; + for (int32_t i = i_draft_beg; i < n_block_tokens; ++i) { const int32_t idx = beg + i; if (conf && conf[(size_t) idx * n_embd_dec] < params.p_min) { diff --git a/conversion/__init__.py b/conversion/__init__.py index 4731b62953a9..3232a1050bfe 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -55,6 +55,8 @@ "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", "Qwen3DSparkModel": "qwen", + "DSparkDraftModel": "qwen", + "DSparkSpeculator": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DeepseekV4DSparkModel": "deepseek", "DistilBertForMaskedLM": "bert", diff --git a/conversion/qwen.py b/conversion/qwen.py index 7a1c11317b35..26b10452b6f9 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -4,12 +4,13 @@ from typing import Any, Callable, Iterable, TYPE_CHECKING +import numpy as np import torch if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, gguf, logger +from .base import LazyTorchTensor, ModelBase, TextModel, gguf, logger @ModelBase.register("QWenLMHeadModel") @@ -708,22 +709,82 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("Qwen3DSparkModel") +@ModelBase.register("Qwen3DSparkModel", "DSparkDraftModel", "DSparkSpeculator") @ModelBase.example("satgeze/Qwen3.6-27B-DSpark") class DSparkModel(DFlashModel): - # DSpark = DFlash + a semi-autoregressive Markov head + # DSpark = DFlash + a semi-autoregressive Markov head. model_arch = gguf.MODEL_ARCH.DFLASH - 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 - }) + def __init__(self, dir_model, *args, **kwargs): + hparams = kwargs.pop("hparams", None) + if hparams is None: + hparams = ModelBase.load_hparams(dir_model, False) + + # EAGLE3-style exports use the 1+N bonus-anchor block, DFlash-lineage exports sample from the anchor + self._sample_from_anchor = hparams.get( + "sample_from_anchor", + "transformer_layer_config" not in hparams and "aux_hidden_state_layer_ids" not in hparams) + if "transformer_layer_config" in hparams: + hparams = {**hparams, **hparams["transformer_layer_config"]} + + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + # normalize both schemas to DFlash's nested dflash_config + if "aux_hidden_state_layer_ids" in self.hparams: + self.hparams.setdefault("dflash_config", { + "mask_token_id": self.hparams.get("mask_token_id"), + "target_layer_ids": [i - 1 for i in self.hparams["aux_hidden_state_layer_ids"]], + }) + else: + self.hparams.setdefault("dflash_config", { + k: self.hparams[k] for k in ("target_layer_ids", "mask_token_id") if k in self.hparams + }) + + if (markov_head_type := self.hparams.get("markov_head_type", "vanilla")) != "vanilla": + raise ValueError(f"unsupported markov_head_type {markov_head_type!r} (only 'vanilla' is supported)") + + n_vocab = self.hparams["vocab_size"] + self._n_vocab_draft = self.hparams.get("draft_vocab_size") or n_vocab + if self._n_vocab_draft > n_vocab: + raise ValueError(f"draft_vocab_size {self._n_vocab_draft} exceeds vocab_size {n_vocab}") + self._d2t: Tensor | None = None + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: - name, gen = item - if name.endswith(("embed_tokens.weight", "lm_head.weight")): + if item[0] == "t2d": # not used at runtime return None - return super().filter_tensors((name, gen)) + return super().filter_tensors(item) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name == "model.d2t": + self._d2t = data_torch + return + + if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")): + return + + yield from super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + + n_vocab = self.hparams["vocab_size"] + if self._n_vocab_draft < n_vocab and self._d2t is None: + raise ValueError(f"draft_vocab_size {self._n_vocab_draft} < vocab_size {n_vocab} but no d2t table found") + + # write d2t as absolute target token ids + if self._d2t is not None: + data = LazyTorchTensor.to_eager(self._d2t).to(torch.int64).cpu().numpy().reshape(-1) + if data.size != self._n_vocab_draft: + raise ValueError(f"d2t size {data.size} does not match draft_vocab_size {self._n_vocab_draft}") + data = data + np.arange(data.size, dtype=np.int64) + if np.any((data < 0) | (data >= n_vocab)): + raise ValueError(f"d2t target ids out of range for target vocab size {n_vocab}") + if np.unique(data).size != data.size: + raise ValueError("d2t contains duplicate target ids") + logger.info(f"{'d2t,':<30} --> I64, shape = {{{data.size}}}") + self.gguf_writer.add_tensor("d2t", data, raw_dtype=gguf.GGMLQuantizationType.I64) diff --git a/docs/speculative.md b/docs/speculative.md index 25abef1b602b..0f9f8a3d977a 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -106,6 +106,10 @@ acceptance (from the draft's confidence head, if present) falls below `P` (defau Currently only drafts with a Qwen3 backbone are supported; support for other backbones (e.g. Gemma4) is planned. +DSpark drafts exported in the [speculators](https://github.com/vllm-project/speculators) format +(for example [`RedHatAI/gemma-4-31B-it-speculator.dspark`](https://huggingface.co/RedHatAI/gemma-4-31B-it-speculator.dspark)) +convert the same way. + See: - #25173 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 74a1a2bd357a..d043c9b6ecce 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -162,6 +162,7 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" @@ -4819,6 +4820,7 @@ class MODEL_TENSOR(IntEnum): ], MODEL_ARCH.DFLASH: [ MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.ATTN_Q, @@ -4858,6 +4860,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_SHEXP, MODEL_TENSOR.FC, MODEL_TENSOR.ENC_OUTPUT_NORM, + MODEL_TENSOR.D2T, # optional DSpark heads MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 475008411c04..9e0914fd866b 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -981,6 +981,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_sample_from_anchor(self, value: bool) -> None: + self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.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 b23bc69153d2..3292942b4133 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -76,14 +76,14 @@ class TensorNameMap: # Output MODEL_TENSOR.OUTPUT: ( "embed_out", # gptneox - "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2 + "lm_head", # gpt2 mpt falcon llama-hf baichuan qwen mamba dbrx jais nemotron exaone olmoe olmo2 phimoe plamo2 llama4 "output", # llama-pth bloom internlm2 "word_embeddings_for_head", # persimmon "lm_head.linear", # phi2 "output_layer", # chatglm "head", # rwkv "head.out", # wavtokenizer - "lm_head", # llama4 + "model.lm_head", # dflash "model.transformer.ff_out", # llada "head.decoder", # modern-bert ), diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index 39e4f5dd8c51..5b70a5179496 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -85,6 +85,16 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t n_embd_inp = hparams.n_embd_inp_enc(); tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + + // reduced draft vocab (optional): d2t maps draft rows to target token ids + int64_t n_vocab_draft = n_vocab; + const struct ggml_tensor * d2t_meta = ml->get_tensor_meta("d2t"); + if (d2t_meta) { + n_vocab_draft = d2t_meta->ne[0]; + d2t = create_tensor(tn(LLM_TENSOR_D2T), { n_vocab_draft }, 0); + LLAMA_LOG_INFO("%s: DFlash using d2t mapping (draft_vocab_size = %lld)\n", __func__, (long long) n_vocab_draft); + } + // DSpark = DFlash + a semi-autoregressive Markov head and Confidence head // // TODO: only Qwen3-style backbones are supported for now; other backbones (e.g. Gemma4) @@ -94,7 +104,7 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { const int64_t dspark_markov_rank = markov_meta->ne[0]; dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); @@ -157,6 +167,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { return; } + // optional: reduced-vocab drafts ship their own, full-vocab drafts share the target's via ctx_other + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED); + for (int i = 0; i < n_layer; ++i) { auto & layer = layers[i]; @@ -242,6 +255,11 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & const int64_t block_size = std::stoi(it->second); GGML_ASSERT(block_size > 0); + // bonus anchor (SpecForge exports): slot 0 is a bonus token, not a prediction slot + const auto it_anchor = model.gguf_kv.find("dflash.sample_from_anchor"); + const bool sample_from_anchor = it_anchor == model.gguf_kv.end() || it_anchor->second == "true"; + const int64_t i_draft_beg = sample_from_anchor ? 0 : 1; + const int64_t n_blocks = g.ubatch.n_seqs_unq; GGML_ASSERT(n_blocks > 0 && n_tok % n_blocks == 0 && "DSpark markov head requires equal-size blocks"); // runtime tokens per block in this ubatch (anchor + drafted positions), bounded by training block_size @@ -263,11 +281,26 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * cat = nullptr; ggml_tensor * cat_conf = nullptr; + if (!sample_from_anchor) { + // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column + cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + } + // 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) { + for (int64_t i = i_draft_beg; 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] + ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks] + if (model.d2t) { + // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others) + const int64_t n_draft_vocab = bias->ne[0]; + ggml_tensor * full = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_blocks), 0.0f); + bias = ggml_set_rows(ctx0, full, + ggml_reshape_3d(ctx0, bias, 1, n_draft_vocab, n_blocks), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + bias = ggml_reshape_2d(ctx0, bias, n_vocab, n_blocks); + } // position i of every block: strided view [n_vocab, n_blocks] ggml_tensor * base_i = ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, i*base->nb[1]); @@ -497,6 +530,22 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra } cur = build_lora_mm(output, cur, output_s); + + // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t + if (model.d2t) { + const int64_t n_draft_vocab = cur->ne[0]; + const int64_t n_outputs = cur->ne[1]; + const int64_t n_vocab = (int64_t) model.vocab.n_tokens(); + + GGML_ASSERT(model.d2t->type == GGML_TYPE_I64); + GGML_ASSERT(model.d2t->ne[0] == n_draft_vocab); + + ggml_tensor * logits = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_vocab, n_outputs), -INFINITY); + cur = ggml_set_rows(ctx0, logits, + ggml_reshape_3d(ctx0, cur, 1, n_draft_vocab, n_outputs), + ggml_reshape_3d(ctx0, model.d2t, n_draft_vocab, 1, 1)); + cur = ggml_reshape_2d(ctx0, cur, n_vocab, n_outputs); + } cb(cur, "result_output", -1); res->t_logits = cur;