[Spec Decode] DSpark speculators checkpoint support (#47093)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Jul 1, 2026 at 20:32 UTC
2b753ad200d52a2dc16e61ff3c92a45711e2750c
4 files changed
+140
-16
vllm/model_executor/models/qwen3_dspark.py
+37
-5
@@ -36,18 +36,26 @@ logger = init_logger(__name__)
36
class DSparkMarkovHead(nn.Module):
37
"""Sequential transition-bias head (low-rank V x r, r x V).
38
39
- ``markov_w1[token]`` is an r-dim embedding of the previously sampled token;
40
- ``markov_w2`` projects it back to a vocab-size bias added to the base logits.
39
+ ``markov_w1[token]`` embeds the previously sampled token (target vocab,
40
+ ``vocab_size``); ``markov_w2`` projects it to a draft-vocab bias
41
+ (``draft_vocab_size``) added to the base draft logits. The two sizes
42
+ coincide for full-vocab drafts.
43
"""
44
43
- def __init__(self, vocab_size: int, markov_rank: int, prefix: str) -> None:
45
+ def __init__(
46
+ self,
47
+ vocab_size: int,
48
+ draft_vocab_size: int,
49
+ markov_rank: int,
50
+ prefix: str,
51
+ ) -> None:
52
super().__init__()
53
# TODO(ben): profile for which (if any) it makes sense to replicate or TP-shard
54
self.markov_w1 = VocabParallelEmbedding(
55
vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w1")
56
)
57
self.markov_w2 = ParallelLMHead(
50
- vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
58
+ draft_vocab_size, markov_rank, prefix=maybe_prefix(prefix, "markov_w2")
59
)
60
61
def embed(self, token_ids: torch.Tensor) -> torch.Tensor:
@@ -73,8 +81,12 @@ class Qwen3DSparkModel(DFlashQwen3Model):
81
vllm_config=vllm_config, start_layer_id=start_layer_id, prefix=prefix
82
)
83
config = self.config
84
+ draft_vocab_size = (
85
+ getattr(config, "draft_vocab_size", None) or config.vocab_size
86
+ )
87
self.markov_head = DSparkMarkovHead(
88
config.vocab_size,
89
+ draft_vocab_size,
90
config.markov_rank,
91
prefix=maybe_prefix(prefix, "markov_head"),
92
)
@@ -117,6 +129,17 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
129
def get_draft_kv_cache_layer_names(self) -> list[str]:
130
return [layer.self_attn.attn.layer_name for layer in self.model.layers]
131
132
+ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
133
+ # Draft-vocab logits without the d2t scatter: the speculator adds the
134
+ # Markov bias in draft space, then remaps via map_draft_to_target.
135
+ return self.logits_processor(self.lm_head, hidden_states)
136
+
137
+ def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
138
+ # Map draft-vocab ids to target ids (identity for full-vocab drafts).
139
+ if self.draft_id_to_target_id is None:
140
+ return draft_ids
141
+ return draft_ids + self.draft_id_to_target_id[draft_ids]
142
+
143
def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
144
return self.model.markov_head.embed(token_ids)
145
@@ -127,8 +150,15 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
150
model_weights = {}
151
includes_embed_tokens = False
152
includes_lm_head = False
153
+ includes_draft_id_mapping = False
154
for name, loaded_weight in weights:
131
- if "lm_head" not in name:
155
+ # t2d is training-only; the draft remaps via d2t at sampling time.
156
+ if "t2d" in name:
157
+ continue
158
+ if "d2t" in name:
159
+ name = name.replace("d2t", "draft_id_to_target_id")
160
+ includes_draft_id_mapping = True
161
+ elif "lm_head" not in name:
162
name = "model." + name
163
if "embed_tokens" in name:
164
includes_embed_tokens = True
@@ -148,6 +178,8 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
178
skip_substrs.append("embed_tokens")
179
if not includes_lm_head:
180
skip_substrs.append("lm_head")
181
+ if not includes_draft_id_mapping:
182
+ skip_substrs.append("draft_id_to_target_id")
183
loader = AutoWeightsLoader(self, skip_substrs=skip_substrs)
184
loader.load_weights(model_weights.items())
185
self.model._build_fused_kv_buffers()
vllm/models/deepseek_v4/nvidia/dspark.py
+11
@@ -110,8 +110,12 @@ class DSparkDeepseekV4Model(nn.Module):
110
self.hc_head_scale = nn.Parameter(
111
torch.empty(1, dtype=torch.float32), requires_grad=False
112
)
113
+ draft_vocab_size = (
114
+ getattr(config, "draft_vocab_size", None) or config.vocab_size
115
+ )
116
self.markov_head = DSparkMarkovHead(
117
config.vocab_size,
118
+ draft_vocab_size,
119
config.dspark_markov_rank,
120
prefix=maybe_prefix(prefix, "markov_head"),
121
)
@@ -318,6 +322,13 @@ class DSparkDeepseekV4ForCausalLM(nn.Module):
322
"""Base logits U_k = lm_head(norm(head_hidden))."""
323
return self.logits_processor(self.lm_head, self.model.norm(hidden_states))
324
325
+ def compute_draft_logits(self, hidden_states: torch.Tensor) -> torch.Tensor:
326
+ # Full-vocab draft: base logits, no d2t scatter.
327
+ return self.compute_logits(hidden_states)
328
+
329
+ def map_draft_to_target(self, draft_ids: torch.Tensor) -> torch.Tensor:
330
+ return draft_ids # full-vocab: draft ids are target ids
331
+
332
def markov_embed(self, token_ids: torch.Tensor) -> torch.Tensor:
333
return self.model.markov_head.embed(token_ids)
334
vllm/transformers_utils/configs/speculators/algos.py
+44
@@ -120,3 +120,47 @@ def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
120
"mask_token_id": config_dict["mask_token_id"],
121
"target_layer_ids": [i - 1 for i in aux_layer_ids],
122
}
123
+
124
+
125
+@register_speculator("dspark")
126
+def update_dspark(config_dict: dict, pre_trained_config: dict) -> None:
127
+ """
128
+ Apply DSpark specific configuration transformations to the `dict` used to
129
+ construct the Transformers PreTrainedConfig.
130
+
131
+ DSpark extends DFlash with a Markov logit-bias head, reusing the same
132
+ Qwen3DSparkModel loader and DSparkSpeculator runtime as the dense DSpark
133
+ checkpoints (e.g. deepseek-ai/dspark_qwen3_8b_block7).
134
+
135
+ DSpark specific fields:
136
+ - draft_vocab_size: draft vocab size; when smaller than the target vocab the
137
+ checkpoint also ships d2t/t2d remap tables.
138
+ - mask_token_id (required): token id for parallel-drafting mask slots.
139
+ - markov_rank / markov_head_type: low-rank Markov logit-bias head.
140
+ - block_size: semi-autoregressive draft block size.
141
+ - enable_confidence_head / confidence_head_with_markov: confidence head.
142
+ - aux_hidden_state_layer_ids (required): target layer indices feeding the
143
+ drafter. Mapped to both eagle_aux_hidden_state_layer_ids and
144
+ target_layer_ids (DSpark's i-1 layer semantics).
145
+ """
146
+ pre_trained_config["architectures"] = ["Qwen3DSparkModel"]
147
+ # Speculators DSpark uses the 1+N fill-in block (anchor is a bonus token).
148
+ pre_trained_config["dspark_bonus_anchor"] = True
149
+
150
+ aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
151
+ pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
152
+ # DSpark indexes target layers as aux_id - 1 (matches the dense configs).
153
+ pre_trained_config["target_layer_ids"] = [i - 1 for i in aux_layer_ids]
154
+
155
+ for key in (
156
+ "draft_vocab_size",
157
+ "target_hidden_size",
158
+ "mask_token_id",
159
+ "markov_rank",
160
+ "markov_head_type",
161
+ "block_size",
162
+ "enable_confidence_head",
163
+ "confidence_head_with_markov",
164
+ ):
165
+ if config_dict.get(key) is not None:
166
+ pre_trained_config[key] = config_dict[key]
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+48
-11
@@ -13,6 +13,8 @@ Differences from DFlash:
13
token), so we sample at all N positions and ``sample_pos = query_pos + 1``
14
(standard next-token), whereas DFlash's masks sit AT the predicted position.
15
This is the ``sample_from_anchor`` path in the shared prepare-inputs kernel.
16
+ Speculators-format checkpoints instead use the DFlash ``1 + N`` fill-in
17
+ layout (anchor is the bonus token).
18
* Sequential Markov sampling: instead of DFlash's single parallel sample, we
19
sample left-to-right, adding a prefix-dependent Markov bias derived from the
20
previously sampled token at each step.
@@ -38,8 +40,15 @@ class DSparkSpeculator(DFlashSpeculator):
40
def __init__(self, vllm_config: VllmConfig, device: torch.device):
41
super().__init__(vllm_config, device)
42
41
- # Anchor-first: N query tokens per request (anchor + N-1 noise), not 1+N.
42
- self.num_query_per_req = self.num_speculative_steps
43
+ # Anchor-as-first (N slots) unless the checkpoint uses the 1+N fill-in
44
+ # block, where the anchor is a separate bonus token.
45
+ self.sample_from_anchor = not getattr(
46
+ self.draft_model_config.hf_config, "dspark_bonus_anchor", False
47
+ )
48
+ if self.sample_from_anchor:
49
+ self.num_query_per_req = self.num_speculative_steps
50
+ else:
51
+ self.num_query_per_req = 1 + self.num_speculative_steps
52
53
# DSpark consumes mean-pooled target aux hidden states at the target
54
# layers, combined to hidden_size via main_proj. Store that combined
@@ -52,9 +61,6 @@ class DSparkSpeculator(DFlashSpeculator):
61
62
self.dflash_causal = False
63
55
- # The anchor query position is itself a prediction (see module docstring).
56
- self.sample_from_anchor = True
57
-
64
self._step_cols = torch.arange(
65
self.num_speculative_steps, dtype=torch.int32, device=device
66
)
@@ -64,12 +70,33 @@ class DSparkSpeculator(DFlashSpeculator):
70
* self.num_query_per_req
71
)
72
73
+ # Reduced-vocab probabilistic drafting only; set in load_draft_model.
74
+ self._d2t_scatter_index: torch.Tensor | None = None
75
+ self._draft_scatter_buf: torch.Tensor | None = None
76
+
77
def load_draft_model(
78
self,
79
target_model: torch.nn.Module,
80
target_attn_layer_names: set[str],
81
) -> torch.nn.Module:
72
- return load_dspark_model(target_model, self.vllm_config)
82
+ model = load_dspark_model(target_model, self.vllm_config)
83
+ # Reduced draft vocab: probabilistic rejection sampling indexes draft
84
+ # logits by target id, so precompute the draft->target column map and a
85
+ # scratch buffer to scatter logits into target vocab before sampling.
86
+ if self.draft_logits is not None and model.draft_id_to_target_id is not None:
87
+ d2t = model.draft_id_to_target_id
88
+ self._d2t_scatter_index = (
89
+ torch.arange(d2t.shape[0], device=d2t.device) + d2t
90
+ )
91
+ # -inf once; the per-step scatter overwrites the draft->target
92
+ # columns. Kept separate from draft_logits to avoid aliasing.
93
+ self._draft_scatter_buf = torch.full(
94
+ (self.max_num_reqs, self.vocab_size),
95
+ float("-inf"),
96
+ dtype=self.draft_logits.dtype,
97
+ device=self.device,
98
+ )
99
+ return model
100
101
def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
102
# Sequential Markov sampling over the backbone's output hidden states.
@@ -77,7 +104,8 @@ class DSparkSpeculator(DFlashSpeculator):
104
num_sample = num_reqs * n_spec
105
# Per-(req, position) head hidden, ordered (req, step).
106
sample_hidden = head_hidden[self.sample_indices[:num_sample]]
80
- base_logits = self.model.compute_logits(sample_hidden)
107
+ # Draft-vocab logits; sampled ids are remapped to target vocab below.
108
+ base_logits = self.model.compute_draft_logits(sample_hidden)
109
vocab_size = base_logits.shape[-1]
110
base_logits = base_logits.view(num_reqs, n_spec, vocab_size)
111
@@ -94,9 +122,16 @@ class DSparkSpeculator(DFlashSpeculator):
122
bias = self.model.markov_bias(markov_embed)
123
logits_i = base_logits[:, i] + bias
124
if self.draft_logits is not None:
125
+ # Probabilistic: sample in target vocab (a reduced draft vocab is
126
+ # scattered into its target columns; full vocab is already there).
127
+ if self._d2t_scatter_index is not None:
128
+ assert self._draft_scatter_buf is not None
129
+ buf = self._draft_scatter_buf[:num_reqs]
130
+ buf.index_copy_(1, self._d2t_scatter_index, logits_i.to(buf.dtype))
131
+ logits_i = buf
132
# sample_pos is the predicted token's position Q; the target
133
# verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1.
99
- draft_i = gumbel_sample(
134
+ draft_sampled_i = gumbel_sample(
135
logits_i,
136
idx_map[:, i],
137
self.temperature,
@@ -108,9 +143,11 @@ class DSparkSpeculator(DFlashSpeculator):
143
use_fp64=self.use_fp64_gumbel,
144
)
145
else:
111
- draft_i = logits_i.argmax(dim=-1)
112
- self.draft_tokens[:num_reqs, i] = draft_i
113
- prev = draft_i
146
+ draft_sampled_i = self.model.map_draft_to_target(
147
+ logits_i.argmax(dim=-1)
148
+ )
149
+ self.draft_tokens[:num_reqs, i] = draft_sampled_i
150
+ prev = draft_sampled_i
151
152
def _generate_draft(
153
self,