[Spec Decode] Add top-k DSpark Markov projection (#49969)
Signed-off-by: Andrii Skliar <askliar@nvidia.com> Co-authored-by: Andrii Skliar <askliar@nvidia.com>
Andrii Skliar committed
Aug 4, 2026 at 14:52 UTC
5789897aa40fbab6bdfcffaa9e83da64939286fb
4 files changed
+235
-26
tests/v1/spec_decode/test_dspark_topk.py
new
+58
@@ -0,0 +1,58 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch import nn
6
+
7
+from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead
8
+
9
+
10
+def _markov_head(weight: torch.Tensor) -> DSparkMarkovHead:
11
+ head = DSparkMarkovHead.__new__(DSparkMarkovHead)
12
+ nn.Module.__init__(head)
13
+ head.markov_w2 = nn.Linear(
14
+ weight.shape[1], weight.shape[0], bias=False, dtype=weight.dtype
15
+ )
16
+ head.markov_w2.weight.data.copy_(weight)
17
+ return head
18
+
19
+
20
+def test_gathered_markov_bias_overwrites_dense_logits():
21
+ weight = torch.arange(21, dtype=torch.float32).view(7, 3) / 10
22
+ markov_embed = torch.tensor([[0.5, -1.0, 0.25], [1.0, 0.5, -0.5]])
23
+ logits = torch.tensor(
24
+ [
25
+ [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7],
26
+ [0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1],
27
+ ]
28
+ )
29
+ values, index = logits.topk(3, dim=-1)
30
+ values = torch.stack((values, torch.zeros_like(values)), dim=1)[:, 0]
31
+ expected = values + torch.bmm(weight[index], markov_embed.unsqueeze(-1)).squeeze(-1)
32
+ logits.fill_(float("-inf"))
33
+
34
+ result = _markov_head(weight).apply_bias_gathered(
35
+ markov_embed, logits, values, index
36
+ )
37
+
38
+ assert result is logits
39
+ torch.testing.assert_close(result.gather(1, index), expected)
40
+ selected = torch.zeros_like(result, dtype=torch.bool).scatter_(1, index, True)
41
+ assert torch.isneginf(result.masked_select(~selected)).all()
42
+
43
+
44
+def test_gathered_markov_bias_matches_dense_at_full_vocab():
45
+ weight = torch.arange(15, dtype=torch.float32).view(5, 3) / 10
46
+ markov_embed = torch.tensor([[0.5, -1.0, 0.25]])
47
+ logits = torch.tensor([[0.1, 0.4, -0.2, 0.3, 0.0]])
48
+ original = logits.clone()
49
+ values, index = logits.topk(logits.shape[-1], dim=-1)
50
+ scale = 0.5
51
+ logits.fill_(float("-inf"))
52
+
53
+ result = _markov_head(weight).apply_bias_gathered(
54
+ markov_embed, logits, values, index, scale
55
+ )
56
+
57
+ expected = original + markov_embed @ weight.T * scale
58
+ torch.testing.assert_close(result, expected)
vllm/config/speculative.py
+59
-1
@@ -289,6 +289,10 @@ class SpeculativeConfig:
289
during rejection sampling. This comes at the cost of additional GPU memory
290
usage."""
291
292
+ dspark_draft_topk: int | None = Field(default=None, ge=1)
293
+ """For Qwen3 DSpark drafting, evaluate the Markov projection only for the
294
+ top-k base-logit candidates. Requires draft tensor parallel size 1."""
295
+
296
def compute_hash(self) -> str:
297
"""
298
WARNING: Whenever a new field is added to this config,
@@ -1032,6 +1036,61 @@ class SpeculativeConfig:
1036
"Inkling MTP currently supports exactly one speculative token"
1037
)
1038
1039
+ if self.dspark_draft_topk is not None and self.method != "dspark":
1040
+ raise ValueError("dspark_draft_topk is only supported by DSpark")
1041
+
1042
+ dspark_draft_topk = None
1043
+ if self.method == "dspark":
1044
+ # DSpark is a semi-autoregressive *block* drafter. A
1045
+ # speculative length smaller than the checkpoint's block
1046
+ # feeds the block / Markov-head machinery an unsupported
1047
+ # layout and yields incorrect (garbled) output rather than
1048
+ # merely lower acceptance. Require num_speculative_tokens to
1049
+ # be at least the block size (e.g. 5 or 7 for DeepSeek-V4).
1050
+ dspark_block_size = getattr(
1051
+ self.draft_model_config.hf_config,
1052
+ "dspark_block_size",
1053
+ None,
1054
+ )
1055
+ if (
1056
+ dspark_block_size is not None
1057
+ and self.num_speculative_tokens < dspark_block_size
1058
+ ):
1059
+ raise ValueError(
1060
+ "DSpark requires num_speculative_tokens >= "
1061
+ f"dspark_block_size ({dspark_block_size}); got "
1062
+ f"{self.num_speculative_tokens}. Smaller values "
1063
+ "produce incorrect output. Use "
1064
+ f"num_speculative_tokens={dspark_block_size} or "
1065
+ "larger (e.g. 7)."
1066
+ )
1067
+
1068
+ hf_config = self.draft_model_config.hf_config
1069
+ dspark_draft_topk = self.dspark_draft_topk
1070
+ if dspark_draft_topk is None:
1071
+ dspark_draft_topk = getattr(
1072
+ hf_config, "dspark_draft_topk", None
1073
+ )
1074
+ if dspark_draft_topk is not None:
1075
+ draft_vocab_size = (
1076
+ getattr(hf_config, "draft_vocab_size", None)
1077
+ or hf_config.vocab_size
1078
+ )
1079
+ if not 1 <= dspark_draft_topk <= draft_vocab_size:
1080
+ raise ValueError(
1081
+ "dspark_draft_topk must be between 1 and the "
1082
+ f"draft vocabulary size ({draft_vocab_size})"
1083
+ )
1084
+ if (
1085
+ "Qwen3DSparkModel"
1086
+ not in self.draft_model_config.architectures
1087
+ ):
1088
+ raise ValueError(
1089
+ "dspark_draft_topk is only supported by "
1090
+ "Qwen3DSparkModel"
1091
+ )
1092
+ hf_config.dspark_draft_topk = dspark_draft_topk
1093
+
1094
self.draft_tensor_parallel_size = (
1095
SpeculativeConfig._verify_and_get_draft_tp(
1096
self.target_parallel_config,
@@ -1039,7 +1098,6 @@ class SpeculativeConfig:
1098
self.draft_model_config.hf_config,
1099
)
1100
)
1042
-
1101
self.draft_model_config.max_model_len = (
1102
SpeculativeConfig._maybe_override_draft_max_model_len(
1103
self.max_model_len,
vllm/model_executor/models/qwen3_dspark.py
+39
@@ -77,6 +77,30 @@ class DSparkMarkovHead(nn.Module):
77
"""Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V])."""
78
return logits_processor(self.markov_w2, markov_embed)
79
80
+ def apply_bias_gathered(
81
+ self,
82
+ markov_embed: torch.Tensor,
83
+ logits: torch.Tensor,
84
+ values: torch.Tensor,
85
+ index: torch.Tensor,
86
+ scale: float = 1.0,
87
+ ) -> torch.Tensor:
88
+ """Apply the Markov bias only to selected rows of ``logits``.
89
+
90
+ The caller initializes ``logits`` to ``-inf`` once for all draft
91
+ positions. This method scatters the corrected candidate values into
92
+ that dense buffer so the normal sampler sees the truncated proposal.
93
+ """
94
+ weight = self.markov_w2.weight[index]
95
+ corrected = values.unsqueeze(-1)
96
+ corrected.baddbmm_(
97
+ weight,
98
+ markov_embed.unsqueeze(-1),
99
+ beta=1.0,
100
+ alpha=scale,
101
+ )
102
+ return logits.scatter_(1, index, corrected.squeeze(-1))
103
+
104
105
class Qwen3DSparkModel(DFlashQwen3Model):
106
"""DFlash Qwen3 backbone + DSpark Markov head."""
@@ -158,6 +182,21 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM):
182
def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor:
183
return self.model.markov_head.bias(markov_embed, self.logits_processor)
184
185
+ def apply_markov_bias_gathered(
186
+ self,
187
+ markov_embed: torch.Tensor,
188
+ logits: torch.Tensor,
189
+ values: torch.Tensor,
190
+ index: torch.Tensor,
191
+ ) -> torch.Tensor:
192
+ return self.model.markov_head.apply_bias_gathered(
193
+ markov_embed,
194
+ logits,
195
+ values,
196
+ index,
197
+ self.logits_processor.scale,
198
+ )
199
+
200
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
201
model_weights = {}
202
includes_embed_tokens = False
vllm/v1/worker/gpu/spec_decode/dspark/speculator.py
+79
-25
@@ -72,6 +72,9 @@ class DSparkSpeculator(DFlashSpeculator):
72
# Reduced-vocab probabilistic drafting only; set in load_draft_model.
73
self._d2t_scatter_index: torch.Tensor | None = None
74
self._draft_scatter_buf: torch.Tensor | None = None
75
+ self._draft_topk: int | None = getattr(
76
+ self.draft_model_config.hf_config, "dspark_draft_topk", None
77
+ )
78
79
def load_draft_model(
80
self,
@@ -97,7 +100,43 @@ class DSparkSpeculator(DFlashSpeculator):
100
)
101
return model
102
103
+ def _sample_logits(
104
+ self,
105
+ logits: torch.Tensor,
106
+ idx_map: torch.Tensor,
107
+ sample_pos: torch.Tensor,
108
+ step: int,
109
+ ) -> torch.Tensor:
110
+ if self.draft_logits is None:
111
+ return self.model.map_draft_to_target(logits.argmax(dim=-1))
112
+
113
+ # Probabilistic sampling and rejection operate in target-vocabulary
114
+ # space. A reduced draft vocabulary is scattered into its target rows.
115
+ if self._d2t_scatter_index is not None:
116
+ assert self._draft_scatter_buf is not None
117
+ buf = self._draft_scatter_buf[: logits.shape[0]]
118
+ buf.index_copy_(1, self._d2t_scatter_index, logits.to(buf.dtype))
119
+ logits = buf
120
+
121
+ # sample_pos is the predicted token's position Q; the target verifies
122
+ # it with the predecessor's Gumbel key (Q-1). Pass Q-1.
123
+ return gumbel_sample(
124
+ logits,
125
+ idx_map,
126
+ self.temperature,
127
+ self.seeds,
128
+ sample_pos - 1,
129
+ apply_temperature=True,
130
+ output_processed_logits=self.draft_logits,
131
+ output_processed_logits_col=self._step_cols[step],
132
+ use_fp64=self.use_fp64_gumbel,
133
+ )
134
+
135
def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
136
+ if self._draft_topk is not None:
137
+ self._sample_sequential_topk(num_reqs, head_hidden)
138
+ return
139
+
140
# Sequential Markov sampling over the backbone's output hidden states.
141
n_spec = self.num_speculative_steps
142
num_sample = num_reqs * n_spec
@@ -120,31 +159,46 @@ class DSparkSpeculator(DFlashSpeculator):
159
markov_embed = self.model.markov_embed(prev)
160
bias = self.model.markov_bias(markov_embed)
161
logits_i = base_logits[:, i] + bias
123
- if self.draft_logits is not None:
124
- # Probabilistic: sample in target vocab (a reduced draft vocab is
125
- # scattered into its target columns; full vocab is already there).
126
- if self._d2t_scatter_index is not None:
127
- assert self._draft_scatter_buf is not None
128
- buf = self._draft_scatter_buf[:num_reqs]
129
- buf.index_copy_(1, self._d2t_scatter_index, logits_i.to(buf.dtype))
130
- logits_i = buf
131
- # sample_pos is the predicted token's position Q; the target
132
- # verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1.
133
- draft_sampled_i = gumbel_sample(
134
- logits_i,
135
- idx_map[:, i],
136
- self.temperature,
137
- self.seeds,
138
- sample_pos[:, i] - 1,
139
- apply_temperature=True,
140
- output_processed_logits=self.draft_logits,
141
- output_processed_logits_col=self._step_cols[i],
142
- use_fp64=self.use_fp64_gumbel,
143
- )
144
- else:
145
- draft_sampled_i = self.model.map_draft_to_target(
146
- logits_i.argmax(dim=-1)
147
- )
162
+ draft_sampled_i = self._sample_logits(
163
+ logits_i, idx_map[:, i], sample_pos[:, i], i
164
+ )
165
+ self.draft_tokens[:num_reqs, i] = draft_sampled_i
166
+ prev = draft_sampled_i
167
+
168
+ def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None:
169
+ """Apply the sequential Markov head only to top-k base-logit candidates.
170
+
171
+ Candidate selection is done once for all draft positions. At each
172
+ sequential step, the selected logits are corrected in place and every
173
+ other entry is set to ``-inf``. The normal dense sampling and rejection
174
+ paths then consume that truncated distribution unchanged.
175
+ """
176
+ assert self._draft_topk is not None
177
+ n_spec = self.num_speculative_steps
178
+ num_sample = num_reqs * n_spec
179
+ sample_hidden = head_hidden[self.sample_indices[:num_sample]]
180
+ base_logits = self.model.compute_draft_logits(sample_hidden)
181
+ base_logits = base_logits.view(num_reqs, n_spec, -1)
182
+ base_values, draft_indices = base_logits.topk(self._draft_topk, dim=-1)
183
+ # Reuse the dense backbone output as the normal sampler's input. Fill
184
+ # once for all positions, then scatter only the corrected candidates
185
+ # during the sequential loop.
186
+ base_logits.fill_(float("-inf"))
187
+ idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec)
188
+ sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec)
189
+ prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]]
190
+
191
+ for i in range(n_spec):
192
+ markov_embed = self.model.markov_embed(prev)
193
+ logits_i = self.model.apply_markov_bias_gathered(
194
+ markov_embed,
195
+ base_logits[:, i],
196
+ base_values[:, i],
197
+ draft_indices[:, i],
198
+ )
199
+ draft_sampled_i = self._sample_logits(
200
+ logits_i, idx_map[:, i], sample_pos[:, i], i
201
+ )
202
self.draft_tokens[:num_reqs, i] = draft_sampled_i
203
prev = draft_sampled_i
204