| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import VllmConfig |
| 9 | from vllm.config.compilation import CUDAGraphMode |
| 10 | from vllm.triton_utils import tl, triton |
| 11 | from vllm.v1.worker.gpu.sample.gumbel import gumbel_noised_argmax |
| 12 | from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator |
| 13 | |
| 14 | |
| 15 | @triton.jit |
| 16 | def _selector_walk_kernel( |
| 17 | scores_ptr, |
| 18 | candidate_ptr, |
| 19 | sample_pos_ptr, |
| 20 | req_state_ptr, |
| 21 | temperature_ptr, |
| 22 | seeds_ptr, |
| 23 | tokens_ptr, |
| 24 | realized_scores_ptr, |
| 25 | num_steps: tl.constexpr, |
| 26 | top_k: tl.constexpr, |
| 27 | BLOCK_K: tl.constexpr, |
| 28 | SAMPLE_PROBABILISTIC: tl.constexpr, |
| 29 | USE_FP64: tl.constexpr, |
| 30 | ): |
| 31 | row = tl.program_id(0) |
| 32 | offsets = tl.arange(0, BLOCK_K) |
| 33 | mask = offsets < top_k |
| 34 | req_state = tl.load(req_state_ptr + row * num_steps) |
| 35 | valid = req_state >= 0 |
| 36 | temperature = tl.load(temperature_ptr + req_state, mask=valid, other=0.0) |
| 37 | seed = tl.load(seeds_ptr + req_state, mask=valid, other=0) |
| 38 | previous = 0 |
| 39 | for step in range(num_steps): |
| 40 | flat = row * num_steps + step |
| 41 | score_base = (flat * top_k + previous) * top_k |
| 42 | scores = tl.load( |
| 43 | scores_ptr + score_base + offsets, |
| 44 | mask=mask & valid, |
| 45 | other=float("-inf"), |
| 46 | ).to(tl.float64 if USE_FP64 else tl.float32) |
| 47 | candidate_base = flat * top_k |
| 48 | candidates = tl.load( |
| 49 | candidate_ptr + candidate_base + offsets, |
| 50 | mask=mask & valid, |
| 51 | other=0, |
| 52 | ) |
| 53 | |
| 54 | # sample_pos is the predicted token's position P. Sampling keys a draw |
| 55 | # by the position before the sampled token, P-1. |
| 56 | sample_pos = tl.load(sample_pos_ptr + flat) - 1 |
| 57 | _, index = gumbel_noised_argmax( |
| 58 | scores, |
| 59 | candidates, |
| 60 | mask & valid, |
| 61 | seed, |
| 62 | sample_pos, |
| 63 | temperature if SAMPLE_PROBABILISTIC else 0.0, |
| 64 | IS_DRAFTING=True, |
| 65 | USE_FP64=USE_FP64, |
| 66 | ) |
| 67 | |
| 68 | tl.store( |
| 69 | realized_scores_ptr + candidate_base + offsets, |
| 70 | scores, |
| 71 | mask=mask & valid, |
| 72 | ) |
| 73 | token = tl.load(candidate_ptr + candidate_base + index, mask=valid, other=0) |
| 74 | tl.store(tokens_ptr + flat, token, mask=valid) |
| 75 | previous = index |
| 76 | |
| 77 | |
| 78 | @triton.jit |
| 79 | def _cache_draft_logits_kernel( |
| 80 | draft_logits_ptr, |
| 81 | cached_candidate_ptr, |
| 82 | candidate_ptr, |
| 83 | scores_ptr, |
| 84 | req_state_ptr, |
| 85 | draft_logits_stride_0, |
| 86 | draft_logits_stride_1, |
| 87 | num_steps: tl.constexpr, |
| 88 | top_k: tl.constexpr, |
| 89 | BLOCK_K: tl.constexpr, |
| 90 | ): |
| 91 | flat = tl.program_id(0) |
| 92 | req_state = tl.load(req_state_ptr + flat) |
| 93 | step = flat % num_steps |
| 94 | offsets = tl.arange(0, BLOCK_K) |
| 95 | mask = (req_state >= 0) & (offsets < top_k) |
| 96 | candidate_base = flat * top_k |
| 97 | cache_base = (req_state * num_steps + step) * top_k |
| 98 | old_token_ids = tl.load(cached_candidate_ptr + cache_base + offsets, mask=mask) |
| 99 | logits_base = ( |
| 100 | draft_logits_ptr |
| 101 | + req_state * draft_logits_stride_0 |
| 102 | + step * draft_logits_stride_1 |
| 103 | ) |
| 104 | tl.store(logits_base + old_token_ids, -float("inf"), mask=mask) |
| 105 | token_ids = tl.load(candidate_ptr + candidate_base + offsets, mask=mask) |
| 106 | scores = tl.load(scores_ptr + candidate_base + offsets, mask=mask) |
| 107 | tl.store(logits_base + token_ids, scores, mask=mask) |
| 108 | tl.store(cached_candidate_ptr + cache_base + offsets, token_ids, mask=mask) |
| 109 | |
| 110 | |
| 111 | class CandidateSampler: |
| 112 | """The shared DFlash2/LiLiCorr candidate walk and realized proposal cache.""" |
| 113 | |
| 114 | def __init__( |
| 115 | self, max_num_reqs: int, num_steps: int, top_k: int, device: torch.device |
| 116 | ): |
| 117 | self.num_steps = num_steps |
| 118 | self.top_k = top_k |
| 119 | self.scores = torch.empty( |
| 120 | max_num_reqs, num_steps, top_k, dtype=torch.float32, device=device |
| 121 | ) |
| 122 | self.cached_candidate_ids = torch.zeros( |
| 123 | self.scores.shape, dtype=torch.int64, device=device |
| 124 | ) |
| 125 | |
| 126 | def sample( |
| 127 | self, |
| 128 | candidate_ids: torch.Tensor, |
| 129 | scores: torch.Tensor, |
| 130 | num_reqs: int, |
| 131 | sample_pos: torch.Tensor, |
| 132 | idx_mapping: torch.Tensor, |
| 133 | temperature: torch.Tensor, |
| 134 | seeds: torch.Tensor, |
| 135 | draft_tokens: torch.Tensor, |
| 136 | draft_logits: torch.Tensor | None, |
| 137 | use_fp64: bool, |
| 138 | ) -> None: |
| 139 | candidate_ids = candidate_ids.contiguous() |
| 140 | sample_pos = sample_pos.contiguous() |
| 141 | idx_mapping = idx_mapping.contiguous() |
| 142 | block_k = triton.next_power_of_2(self.top_k) |
| 143 | _selector_walk_kernel[(num_reqs,)]( |
| 144 | scores.contiguous(), |
| 145 | candidate_ids.contiguous(), |
| 146 | sample_pos, |
| 147 | idx_mapping, |
| 148 | temperature, |
| 149 | seeds, |
| 150 | draft_tokens, |
| 151 | self.scores, |
| 152 | num_steps=self.num_steps, |
| 153 | top_k=self.top_k, |
| 154 | BLOCK_K=block_k, |
| 155 | SAMPLE_PROBABILISTIC=draft_logits is not None, |
| 156 | USE_FP64=use_fp64, |
| 157 | num_warps=1, |
| 158 | ) |
| 159 | |
| 160 | if draft_logits is not None: |
| 161 | self._cache_draft_logits( |
| 162 | candidate_ids, num_reqs * self.num_steps, idx_mapping, draft_logits |
| 163 | ) |
| 164 | |
| 165 | def _cache_draft_logits( |
| 166 | self, |
| 167 | candidate_ids: torch.Tensor, |
| 168 | num_sample: int, |
| 169 | idx_mapping: torch.Tensor, |
| 170 | draft_logits: torch.Tensor, |
| 171 | ) -> None: |
| 172 | block_k = triton.next_power_of_2(self.top_k) |
| 173 | _cache_draft_logits_kernel[(num_sample,)]( |
| 174 | draft_logits, |
| 175 | self.cached_candidate_ids, |
| 176 | candidate_ids, |
| 177 | self.scores, |
| 178 | idx_mapping, |
| 179 | draft_logits.stride(0), |
| 180 | draft_logits.stride(1), |
| 181 | num_steps=self.num_steps, |
| 182 | top_k=self.top_k, |
| 183 | BLOCK_K=block_k, |
| 184 | num_warps=1, |
| 185 | ) |
| 186 | |
| 187 | |
| 188 | class DFlash2Speculator(DFlashSpeculator): |
| 189 | _speculator_name = "DFlash2" |
| 190 | _candidate_top_k_key = "selector_top_k" |
| 191 | |
| 192 | def __init__(self, vllm_config: VllmConfig, device: torch.device): |
| 193 | super().__init__(vllm_config, device) |
| 194 | draft_config = self.draft_model_config.hf_config.dflash_config |
| 195 | self.top_k = int(draft_config[self._candidate_top_k_key]) |
| 196 | self.candidate_sampler = CandidateSampler( |
| 197 | self.max_num_reqs, self.num_speculative_steps, self.top_k, device |
| 198 | ) |
| 199 | |
| 200 | def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: |
| 201 | # fp32 so the walk and the rejection that checks it read the same |
| 202 | # distribution; -inf because the cache kernel writes only the K |
| 203 | # candidates. |
| 204 | return torch.float32, -float("inf") |
| 205 | |
| 206 | def _generate_draft( |
| 207 | self, |
| 208 | num_reqs: int, |
| 209 | num_tokens_padded: int, |
| 210 | attn_metadata: dict[str, Any] | None, |
| 211 | slot_mappings: dict[str, torch.Tensor] | None, |
| 212 | num_tokens_across_dp: torch.Tensor | None, |
| 213 | cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, |
| 214 | ) -> None: |
| 215 | last_hidden_states = self._run_model( |
| 216 | num_tokens_padded, |
| 217 | attn_metadata, |
| 218 | slot_mappings, |
| 219 | num_tokens_across_dp, |
| 220 | cudagraph_runtime_mode, |
| 221 | ) |
| 222 | num_sample = num_reqs * self.num_speculative_steps |
| 223 | hidden_states = last_hidden_states[self.sample_indices[:num_sample]].view( |
| 224 | num_reqs, self.num_speculative_steps, -1 |
| 225 | ) |
| 226 | candidate_ids, unary_logits = self.model.compute_candidates( |
| 227 | hidden_states.flatten(0, 1) |
| 228 | ) |
| 229 | candidate_ids = candidate_ids.view( |
| 230 | num_reqs, self.num_speculative_steps, self.top_k |
| 231 | ) |
| 232 | unary_logits = unary_logits.view_as(candidate_ids) |
| 233 | scores = self._score_candidates(candidate_ids, unary_logits, hidden_states) |
| 234 | self.candidate_sampler.sample( |
| 235 | candidate_ids, |
| 236 | scores, |
| 237 | num_reqs, |
| 238 | self.sample_pos, |
| 239 | self.sample_idx_mapping, |
| 240 | self.temperature, |
| 241 | self.seeds, |
| 242 | self.draft_tokens, |
| 243 | self.draft_logits, |
| 244 | self.use_fp64_gumbel, |
| 245 | ) |
| 246 | if self.enable_adaptive_verification: |
| 247 | self._maybe_predict_acceptance( |
| 248 | self.candidate_sampler.scores[:num_reqs].flatten(0, 1), |
| 249 | self.sample_idx_mapping[:num_sample], |
| 250 | self.sample_col[:num_sample], |
| 251 | ) |
| 252 | |
| 253 | def _score_candidates( |
| 254 | self, |
| 255 | candidate_ids: torch.Tensor, |
| 256 | unary_logits: torch.Tensor, |
| 257 | hidden_states: torch.Tensor, |
| 258 | ) -> torch.Tensor: |
| 259 | num_reqs = candidate_ids.shape[0] |
| 260 | anchor_token_ids = self.input_buffers.input_ids[ |
| 261 | : num_reqs * self.num_query_per_req : self.num_query_per_req |
| 262 | ] |
| 263 | return self.model.model.candidate_selector( |
| 264 | candidate_ids, unary_logits, hidden_states, anchor_token_ids |
| 265 | ) |