diff --git a/tests/v1/spec_decode/test_dspark_topk.py b/tests/v1/spec_decode/test_dspark_topk.py new file mode 100644 index 00000000000..a33483d6db3 --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_topk.py @@ -0,0 +1,58 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.model_executor.models.qwen3_dspark import DSparkMarkovHead + + +def _markov_head(weight: torch.Tensor) -> DSparkMarkovHead: + head = DSparkMarkovHead.__new__(DSparkMarkovHead) + nn.Module.__init__(head) + head.markov_w2 = nn.Linear( + weight.shape[1], weight.shape[0], bias=False, dtype=weight.dtype + ) + head.markov_w2.weight.data.copy_(weight) + return head + + +def test_gathered_markov_bias_overwrites_dense_logits(): + weight = torch.arange(21, dtype=torch.float32).view(7, 3) / 10 + markov_embed = torch.tensor([[0.5, -1.0, 0.25], [1.0, 0.5, -0.5]]) + logits = torch.tensor( + [ + [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7], + [0.7, 0.6, 0.5, 0.4, 0.3, 0.2, 0.1], + ] + ) + values, index = logits.topk(3, dim=-1) + values = torch.stack((values, torch.zeros_like(values)), dim=1)[:, 0] + expected = values + torch.bmm(weight[index], markov_embed.unsqueeze(-1)).squeeze(-1) + logits.fill_(float("-inf")) + + result = _markov_head(weight).apply_bias_gathered( + markov_embed, logits, values, index + ) + + assert result is logits + torch.testing.assert_close(result.gather(1, index), expected) + selected = torch.zeros_like(result, dtype=torch.bool).scatter_(1, index, True) + assert torch.isneginf(result.masked_select(~selected)).all() + + +def test_gathered_markov_bias_matches_dense_at_full_vocab(): + weight = torch.arange(15, dtype=torch.float32).view(5, 3) / 10 + markov_embed = torch.tensor([[0.5, -1.0, 0.25]]) + logits = torch.tensor([[0.1, 0.4, -0.2, 0.3, 0.0]]) + original = logits.clone() + values, index = logits.topk(logits.shape[-1], dim=-1) + scale = 0.5 + logits.fill_(float("-inf")) + + result = _markov_head(weight).apply_bias_gathered( + markov_embed, logits, values, index, scale + ) + + expected = original + markov_embed @ weight.T * scale + torch.testing.assert_close(result, expected) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index aa3fc67ea11..e6190979684 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -289,6 +289,10 @@ class SpeculativeConfig: during rejection sampling. This comes at the cost of additional GPU memory usage.""" + dspark_draft_topk: int | None = Field(default=None, ge=1) + """For Qwen3 DSpark drafting, evaluate the Markov projection only for the + top-k base-logit candidates. Requires draft tensor parallel size 1.""" + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, @@ -1032,6 +1036,61 @@ class SpeculativeConfig: "Inkling MTP currently supports exactly one speculative token" ) + if self.dspark_draft_topk is not None and self.method != "dspark": + raise ValueError("dspark_draft_topk is only supported by DSpark") + + dspark_draft_topk = None + if self.method == "dspark": + # DSpark is a semi-autoregressive *block* drafter. A + # speculative length smaller than the checkpoint's block + # feeds the block / Markov-head machinery an unsupported + # layout and yields incorrect (garbled) output rather than + # merely lower acceptance. Require num_speculative_tokens to + # be at least the block size (e.g. 5 or 7 for DeepSeek-V4). + dspark_block_size = getattr( + self.draft_model_config.hf_config, + "dspark_block_size", + None, + ) + if ( + dspark_block_size is not None + and self.num_speculative_tokens < dspark_block_size + ): + raise ValueError( + "DSpark requires num_speculative_tokens >= " + f"dspark_block_size ({dspark_block_size}); got " + f"{self.num_speculative_tokens}. Smaller values " + "produce incorrect output. Use " + f"num_speculative_tokens={dspark_block_size} or " + "larger (e.g. 7)." + ) + + hf_config = self.draft_model_config.hf_config + dspark_draft_topk = self.dspark_draft_topk + if dspark_draft_topk is None: + dspark_draft_topk = getattr( + hf_config, "dspark_draft_topk", None + ) + if dspark_draft_topk is not None: + draft_vocab_size = ( + getattr(hf_config, "draft_vocab_size", None) + or hf_config.vocab_size + ) + if not 1 <= dspark_draft_topk <= draft_vocab_size: + raise ValueError( + "dspark_draft_topk must be between 1 and the " + f"draft vocabulary size ({draft_vocab_size})" + ) + if ( + "Qwen3DSparkModel" + not in self.draft_model_config.architectures + ): + raise ValueError( + "dspark_draft_topk is only supported by " + "Qwen3DSparkModel" + ) + hf_config.dspark_draft_topk = dspark_draft_topk + self.draft_tensor_parallel_size = ( SpeculativeConfig._verify_and_get_draft_tp( self.target_parallel_config, @@ -1039,7 +1098,6 @@ class SpeculativeConfig: self.draft_model_config.hf_config, ) ) - self.draft_model_config.max_model_len = ( SpeculativeConfig._maybe_override_draft_max_model_len( self.max_model_len, diff --git a/vllm/model_executor/models/qwen3_dspark.py b/vllm/model_executor/models/qwen3_dspark.py index 04af4566949..03ad01e1d91 100644 --- a/vllm/model_executor/models/qwen3_dspark.py +++ b/vllm/model_executor/models/qwen3_dspark.py @@ -77,6 +77,30 @@ class DSparkMarkovHead(nn.Module): """Vocab-size transition bias from a Markov embedding ([B, r] -> [B, V]).""" return logits_processor(self.markov_w2, markov_embed) + def apply_bias_gathered( + self, + markov_embed: torch.Tensor, + logits: torch.Tensor, + values: torch.Tensor, + index: torch.Tensor, + scale: float = 1.0, + ) -> torch.Tensor: + """Apply the Markov bias only to selected rows of ``logits``. + + The caller initializes ``logits`` to ``-inf`` once for all draft + positions. This method scatters the corrected candidate values into + that dense buffer so the normal sampler sees the truncated proposal. + """ + weight = self.markov_w2.weight[index] + corrected = values.unsqueeze(-1) + corrected.baddbmm_( + weight, + markov_embed.unsqueeze(-1), + beta=1.0, + alpha=scale, + ) + return logits.scatter_(1, index, corrected.squeeze(-1)) + class Qwen3DSparkModel(DFlashQwen3Model): """DFlash Qwen3 backbone + DSpark Markov head.""" @@ -158,6 +182,21 @@ class Qwen3DSparkForCausalLM(DFlashQwen3ForCausalLM): def markov_bias(self, markov_embed: torch.Tensor) -> torch.Tensor: return self.model.markov_head.bias(markov_embed, self.logits_processor) + def apply_markov_bias_gathered( + self, + markov_embed: torch.Tensor, + logits: torch.Tensor, + values: torch.Tensor, + index: torch.Tensor, + ) -> torch.Tensor: + return self.model.markov_head.apply_bias_gathered( + markov_embed, + logits, + values, + index, + self.logits_processor.scale, + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): model_weights = {} includes_embed_tokens = False diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 45dedde3d7d..a3207cd3dad 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -72,6 +72,9 @@ class DSparkSpeculator(DFlashSpeculator): # Reduced-vocab probabilistic drafting only; set in load_draft_model. self._d2t_scatter_index: torch.Tensor | None = None self._draft_scatter_buf: torch.Tensor | None = None + self._draft_topk: int | None = getattr( + self.draft_model_config.hf_config, "dspark_draft_topk", None + ) def load_draft_model( self, @@ -97,7 +100,43 @@ class DSparkSpeculator(DFlashSpeculator): ) return model + def _sample_logits( + self, + logits: torch.Tensor, + idx_map: torch.Tensor, + sample_pos: torch.Tensor, + step: int, + ) -> torch.Tensor: + if self.draft_logits is None: + return self.model.map_draft_to_target(logits.argmax(dim=-1)) + + # Probabilistic sampling and rejection operate in target-vocabulary + # space. A reduced draft vocabulary is scattered into its target rows. + if self._d2t_scatter_index is not None: + assert self._draft_scatter_buf is not None + buf = self._draft_scatter_buf[: logits.shape[0]] + buf.index_copy_(1, self._d2t_scatter_index, logits.to(buf.dtype)) + logits = buf + + # sample_pos is the predicted token's position Q; the target verifies + # it with the predecessor's Gumbel key (Q-1). Pass Q-1. + return gumbel_sample( + logits, + idx_map, + self.temperature, + self.seeds, + sample_pos - 1, + apply_temperature=True, + output_processed_logits=self.draft_logits, + output_processed_logits_col=self._step_cols[step], + use_fp64=self.use_fp64_gumbel, + ) + def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: + if self._draft_topk is not None: + self._sample_sequential_topk(num_reqs, head_hidden) + return + # Sequential Markov sampling over the backbone's output hidden states. n_spec = self.num_speculative_steps num_sample = num_reqs * n_spec @@ -120,31 +159,46 @@ class DSparkSpeculator(DFlashSpeculator): markov_embed = self.model.markov_embed(prev) bias = self.model.markov_bias(markov_embed) logits_i = base_logits[:, i] + bias - if self.draft_logits is not None: - # Probabilistic: sample in target vocab (a reduced draft vocab is - # scattered into its target columns; full vocab is already there). - if self._d2t_scatter_index is not None: - assert self._draft_scatter_buf is not None - buf = self._draft_scatter_buf[:num_reqs] - buf.index_copy_(1, self._d2t_scatter_index, logits_i.to(buf.dtype)) - logits_i = buf - # sample_pos is the predicted token's position Q; the target - # verifies it with the predecessor's Gumbel key (Q-1). Pass Q-1. - draft_sampled_i = gumbel_sample( - logits_i, - idx_map[:, i], - self.temperature, - self.seeds, - sample_pos[:, i] - 1, - apply_temperature=True, - output_processed_logits=self.draft_logits, - output_processed_logits_col=self._step_cols[i], - use_fp64=self.use_fp64_gumbel, - ) - else: - draft_sampled_i = self.model.map_draft_to_target( - logits_i.argmax(dim=-1) - ) + draft_sampled_i = self._sample_logits( + logits_i, idx_map[:, i], sample_pos[:, i], i + ) + self.draft_tokens[:num_reqs, i] = draft_sampled_i + prev = draft_sampled_i + + def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None: + """Apply the sequential Markov head only to top-k base-logit candidates. + + Candidate selection is done once for all draft positions. At each + sequential step, the selected logits are corrected in place and every + other entry is set to ``-inf``. The normal dense sampling and rejection + paths then consume that truncated distribution unchanged. + """ + assert self._draft_topk is not None + n_spec = self.num_speculative_steps + num_sample = num_reqs * n_spec + sample_hidden = head_hidden[self.sample_indices[:num_sample]] + base_logits = self.model.compute_draft_logits(sample_hidden) + base_logits = base_logits.view(num_reqs, n_spec, -1) + base_values, draft_indices = base_logits.topk(self._draft_topk, dim=-1) + # Reuse the dense backbone output as the normal sampler's input. Fill + # once for all positions, then scatter only the corrected candidates + # during the sequential loop. + base_logits.fill_(float("-inf")) + idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) + sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) + prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] + + for i in range(n_spec): + markov_embed = self.model.markov_embed(prev) + logits_i = self.model.apply_markov_bias_gathered( + markov_embed, + base_logits[:, i], + base_values[:, i], + draft_indices[:, i], + ) + draft_sampled_i = self._sample_logits( + logits_i, idx_map[:, i], sample_pos[:, i], i + ) self.draft_tokens[:num_reqs, i] = draft_sampled_i prev = draft_sampled_i