| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DSpark speculator: semi-autoregressive parallel drafting. |
| 4 | |
| 5 | DSpark drafts a block of ``num_speculative_tokens`` tokens in one parallel pass |
| 6 | (reusing the DFlash machinery: context-KV precompute + a query-block forward), |
| 7 | then injects intra-block dependency with a lightweight sequential Markov head. |
| 8 | |
| 9 | Differences from DFlash: |
| 10 | * Anchor-as-first-prediction: each request emits exactly ``N = |
| 11 | num_speculative_tokens`` query tokens (anchor + N-1 noise), NOT ``1 + N``. |
| 12 | Every query position is a prediction (the anchor predicts the first draft |
| 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. |
| 21 | |
| 22 | CUDA graphs (FULL, mirroring DFlash) cover the whole draft step: the parallel |
| 23 | backbone forward AND the sequential Markov sampling. |
| 24 | """ |
| 25 | |
| 26 | from typing import Any |
| 27 | |
| 28 | import torch |
| 29 | |
| 30 | from vllm.config import VllmConfig |
| 31 | from vllm.config.compilation import CUDAGraphMode |
| 32 | from vllm.logger import init_logger |
| 33 | from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample |
| 34 | from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator |
| 35 | from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model |
| 36 | |
| 37 | logger = init_logger(__name__) |
| 38 | |
| 39 | |
| 40 | class DSparkSpeculator(DFlashSpeculator): |
| 41 | _speculator_name = "DSpark" |
| 42 | |
| 43 | def __init__(self, vllm_config: VllmConfig, device: torch.device): |
| 44 | super().__init__(vllm_config, device) |
| 45 | |
| 46 | # Whether to sample from the anchor position. When True, uses anchor-as-first |
| 47 | # (N slots, each position predicts the next token). When False, uses 1+N |
| 48 | # fill-in block (anchor is a bonus token). |
| 49 | self.sample_from_anchor = getattr( |
| 50 | self.draft_model_config.hf_config, "sample_from_anchor", True |
| 51 | ) |
| 52 | if self.sample_from_anchor: |
| 53 | self.num_query_per_req = self.num_speculative_steps |
| 54 | else: |
| 55 | self.num_query_per_req = 1 + self.num_speculative_steps |
| 56 | |
| 57 | # DSpark consumes mean-pooled target aux hidden states at the target |
| 58 | # layers, combined to hidden_size via main_proj. Store that combined |
| 59 | # main_x (hidden_size wide). DSpark does not use the same pre-allocated buffer |
| 60 | # that DeepSeek-V4's MTP uses. |
| 61 | draft_hidden = self.draft_model_config.get_hidden_size() |
| 62 | self.hidden_states = torch.zeros( |
| 63 | self.max_num_tokens, draft_hidden, dtype=self.dtype, device=device |
| 64 | ) |
| 65 | |
| 66 | self._step_cols = torch.arange( |
| 67 | self.num_speculative_steps, dtype=torch.int32, device=device |
| 68 | ) |
| 69 | |
| 70 | self._anchor_idx = ( |
| 71 | torch.arange(self.max_num_reqs, dtype=torch.int64, device=device) |
| 72 | * self.num_query_per_req |
| 73 | ) |
| 74 | |
| 75 | # Reduced-vocab probabilistic drafting only; set in load_draft_model. |
| 76 | self._d2t_scatter_index: torch.Tensor | None = None |
| 77 | self._draft_scatter_buf: torch.Tensor | None = None |
| 78 | self._draft_topk: int | None = getattr( |
| 79 | self.draft_model_config.hf_config, "dspark_draft_topk", None |
| 80 | ) |
| 81 | |
| 82 | self.use_confidence_head: bool = False |
| 83 | |
| 84 | def load_draft_model( |
| 85 | self, |
| 86 | target_model: torch.nn.Module, |
| 87 | target_attn_layer_names: set[str], |
| 88 | ) -> torch.nn.Module: |
| 89 | model = load_dspark_model(target_model, self.vllm_config) |
| 90 | # Reduced draft vocab: probabilistic rejection sampling indexes draft |
| 91 | # logits by target id, so precompute the draft->target column map and a |
| 92 | # scratch buffer to scatter logits into target vocab before sampling. |
| 93 | if self.draft_logits is not None and model.draft_id_to_target_id is not None: |
| 94 | d2t = model.draft_id_to_target_id |
| 95 | self._d2t_scatter_index = ( |
| 96 | torch.arange(d2t.shape[0], device=d2t.device) + d2t |
| 97 | ) |
| 98 | # -inf once; the per-step scatter overwrites the draft->target |
| 99 | # columns. Kept separate from draft_logits to avoid aliasing. |
| 100 | self._draft_scatter_buf = torch.full( |
| 101 | (self.max_num_reqs, self.vocab_size), |
| 102 | float("-inf"), |
| 103 | dtype=self.draft_logits.dtype, |
| 104 | device=self.device, |
| 105 | ) |
| 106 | self.use_confidence_head = ( |
| 107 | self.enable_adaptive_verification |
| 108 | and model.model.confidence_head is not None |
| 109 | ) |
| 110 | if self.use_confidence_head: |
| 111 | # The acceptance estimator is not needed when a trained confidence head |
| 112 | # is available. |
| 113 | self.use_acceptance_estimator = False |
| 114 | return model |
| 115 | |
| 116 | def _sample_logits( |
| 117 | self, |
| 118 | logits: torch.Tensor, |
| 119 | idx_map: torch.Tensor, |
| 120 | sample_pos: torch.Tensor, |
| 121 | step: int, |
| 122 | ) -> torch.Tensor: |
| 123 | self._maybe_predict_acceptance(logits, idx_map, self._step_cols[step]) |
| 124 | if self.draft_logits is None: |
| 125 | draft_ids = logits.argmax(dim=-1) |
| 126 | return self.model.map_draft_to_target(draft_ids) |
| 127 | |
| 128 | # Probabilistic sampling and rejection operate in target-vocabulary |
| 129 | # space. A reduced draft vocabulary is scattered into its target rows. |
| 130 | if self._d2t_scatter_index is not None: |
| 131 | assert self._draft_scatter_buf is not None |
| 132 | buf = self._draft_scatter_buf[: logits.shape[0]] |
| 133 | buf.index_copy_(1, self._d2t_scatter_index, logits.to(buf.dtype)) |
| 134 | logits = buf |
| 135 | |
| 136 | # sample_pos is the predicted token's position P. Sampling keys a draw |
| 137 | # by the position before the sampled token, P-1. |
| 138 | sampler = ( |
| 139 | gumbel_sample |
| 140 | if self.draft_watermarker is None |
| 141 | else self.draft_watermarker.sample |
| 142 | ) |
| 143 | return sampler( |
| 144 | logits, |
| 145 | idx_map, |
| 146 | self.temperature, |
| 147 | self.seeds, |
| 148 | sample_pos - 1, |
| 149 | apply_temperature=True, |
| 150 | is_drafting=True, |
| 151 | logits_cache=self.draft_logits, |
| 152 | logits_cache_col=self._step_cols[step], |
| 153 | use_fp64=self.use_fp64_gumbel, |
| 154 | ) |
| 155 | |
| 156 | def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: |
| 157 | if self._draft_topk is not None: |
| 158 | self._sample_sequential_topk(num_reqs, head_hidden) |
| 159 | return |
| 160 | |
| 161 | # Sequential Markov sampling over the backbone's output hidden states. |
| 162 | n_spec = self.num_speculative_steps |
| 163 | num_sample = num_reqs * n_spec |
| 164 | # Per-(req, position) head hidden, ordered (req, step). |
| 165 | sample_hidden = head_hidden[self.sample_indices[:num_sample]] |
| 166 | # Draft-vocab logits; sampled ids are remapped to target vocab below. |
| 167 | base_logits = self.model.compute_draft_logits(sample_hidden) |
| 168 | vocab_size = base_logits.shape[-1] |
| 169 | base_logits = base_logits.view(num_reqs, n_spec, vocab_size) |
| 170 | |
| 171 | idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) |
| 172 | sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) |
| 173 | confidence_markov_embeds = [] |
| 174 | |
| 175 | # Anchor (bonus) token per request = the input id at query offset 0, |
| 176 | # read via the precomputed persistent index (fixed buffer for capture). |
| 177 | prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] |
| 178 | |
| 179 | for i in range(n_spec): |
| 180 | # Sequential stage: Markov bias from the previously sampled token. |
| 181 | markov_embed = self.model.markov_embed(prev) |
| 182 | if self.use_confidence_head: |
| 183 | confidence_markov_embeds.append(markov_embed) |
| 184 | bias = self.model.markov_bias(markov_embed) |
| 185 | logits_i = base_logits[:, i] + bias |
| 186 | draft_sampled_i = self._sample_logits( |
| 187 | logits_i, idx_map[:, i], sample_pos[:, i], i |
| 188 | ) |
| 189 | self.draft_tokens[:num_reqs, i] = draft_sampled_i |
| 190 | prev = draft_sampled_i |
| 191 | |
| 192 | if self.use_confidence_head: |
| 193 | confidence = self.model.compute_confidence( |
| 194 | sample_hidden, |
| 195 | torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), |
| 196 | ) |
| 197 | self.draft_token_confidence_probs[:num_reqs] = confidence.view( |
| 198 | num_reqs, n_spec |
| 199 | ) |
| 200 | |
| 201 | def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None: |
| 202 | """Apply the sequential Markov head only to top-k base-logit candidates. |
| 203 | |
| 204 | Candidate selection is done once for all draft positions. At each |
| 205 | sequential step, the selected logits are corrected in place and every |
| 206 | other entry is set to ``-inf``. The normal dense sampling and rejection |
| 207 | paths then consume that truncated distribution unchanged. |
| 208 | """ |
| 209 | assert self._draft_topk is not None |
| 210 | n_spec = self.num_speculative_steps |
| 211 | num_sample = num_reqs * n_spec |
| 212 | sample_hidden = head_hidden[self.sample_indices[:num_sample]] |
| 213 | base_logits = self.model.compute_draft_logits(sample_hidden) |
| 214 | base_logits = base_logits.view(num_reqs, n_spec, -1) |
| 215 | base_values, draft_indices = base_logits.topk(self._draft_topk, dim=-1) |
| 216 | # Reuse the dense backbone output as the normal sampler's input. Fill |
| 217 | # once for all positions, then scatter only the corrected candidates |
| 218 | # during the sequential loop. |
| 219 | base_logits.fill_(float("-inf")) |
| 220 | idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) |
| 221 | sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) |
| 222 | confidence_markov_embeds = [] |
| 223 | prev = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] |
| 224 | |
| 225 | for i in range(n_spec): |
| 226 | markov_embed = self.model.markov_embed(prev) |
| 227 | if self.use_confidence_head: |
| 228 | confidence_markov_embeds.append(markov_embed) |
| 229 | logits_i = self.model.apply_markov_bias_gathered( |
| 230 | markov_embed, |
| 231 | base_logits[:, i], |
| 232 | base_values[:, i], |
| 233 | draft_indices[:, i], |
| 234 | ) |
| 235 | draft_sampled_i = self._sample_logits( |
| 236 | logits_i, idx_map[:, i], sample_pos[:, i], i |
| 237 | ) |
| 238 | self.draft_tokens[:num_reqs, i] = draft_sampled_i |
| 239 | prev = draft_sampled_i |
| 240 | |
| 241 | if self.use_confidence_head: |
| 242 | confidence = self.model.compute_confidence( |
| 243 | sample_hidden, |
| 244 | torch.stack(confidence_markov_embeds, dim=1).flatten(0, 1), |
| 245 | ) |
| 246 | self.draft_token_confidence_probs[:num_reqs] = confidence.view( |
| 247 | num_reqs, n_spec |
| 248 | ) |
| 249 | |
| 250 | def _generate_draft( |
| 251 | self, |
| 252 | num_reqs: int, |
| 253 | num_tokens_padded: int, |
| 254 | attn_metadata: dict[str, Any] | None, |
| 255 | slot_mappings: dict[str, torch.Tensor] | None, |
| 256 | num_tokens_across_dp: torch.Tensor | None, |
| 257 | cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, |
| 258 | ) -> None: |
| 259 | # Full draft step (captured under CUDA graph): parallel backbone forward |
| 260 | # then sequential Markov sampling over its hidden state outputs. |
| 261 | head_hidden = self._run_model( |
| 262 | num_tokens_padded, |
| 263 | attn_metadata, |
| 264 | slot_mappings, |
| 265 | num_tokens_across_dp, |
| 266 | cudagraph_runtime_mode, |
| 267 | ) |
| 268 | self._sample_sequential(num_reqs, head_hidden) |