| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections.abc import Iterable, Iterator |
| 4 | from typing import Any |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import SpeculativeConfig |
| 10 | from vllm.config.model import PROCESSED_LOGPROBS_MODES |
| 11 | from vllm.triton_utils import tl, triton |
| 12 | from vllm.v1.outputs import LogprobsTensors |
| 13 | from vllm.v1.spec_decode.utils import unconditional_to_conditional_rates |
| 14 | from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler |
| 15 | from vllm.v1.worker.gpu.input_batch import ( |
| 16 | InputBatch, |
| 17 | get_num_sampled_and_rejected, |
| 18 | ) |
| 19 | from vllm.v1.worker.gpu.metrics.logits import get_num_nans |
| 20 | from vllm.v1.worker.gpu.sample.logprob import compute_topk_scores |
| 21 | from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors |
| 22 | from vllm.v1.worker.gpu.sample.sampler import Sampler |
| 23 | from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS |
| 24 | from vllm.v1.worker.gpu.spec_decode.rejection_sampler_utils import ( |
| 25 | rejection_sample, |
| 26 | ) |
| 27 | |
| 28 | # Cap on the FP32 target-logits buffer materialized by apply_sampling_params. |
| 29 | # TODO(mgoin): Chunking is a workaround. The rejection kernels already upcast |
| 30 | # per vocab block on load and apply ops like temperature and gumbel, so folding |
| 31 | # sampling-param application into those kernels would remove this buffer and |
| 32 | # its traffic entirely. |
| 33 | MAX_CHUNK_BYTES = 2**30 # 1GB |
| 34 | _FP32_BYTES = 4 |
| 35 | |
| 36 | |
| 37 | def get_max_chunk_logits(vocab_size: int) -> int: |
| 38 | """Largest number of logits rows one verification chunk may hold.""" |
| 39 | return max(1, MAX_CHUNK_BYTES // (vocab_size * _FP32_BYTES)) |
| 40 | |
| 41 | |
| 42 | def _iter_request_chunks( |
| 43 | cu_num_logits: np.ndarray, max_chunk_logits: int |
| 44 | ) -> Iterator[tuple[int, int]]: |
| 45 | """Yield maximally packed request ranges without splitting requests.""" |
| 46 | assert max_chunk_logits > 0 |
| 47 | num_reqs = cu_num_logits.size - 1 |
| 48 | start = 0 |
| 49 | while start < num_reqs: |
| 50 | max_logit = int(cu_num_logits[start]) + max_chunk_logits |
| 51 | end = int(np.searchsorted(cu_num_logits, max_logit, side="right") - 1) |
| 52 | end = min(num_reqs, max(start + 1, end)) |
| 53 | yield start, end |
| 54 | start = end |
| 55 | |
| 56 | |
| 57 | @triton.jit |
| 58 | def _flatten_sampled_kernel( |
| 59 | # [num_logits] |
| 60 | flat_sampled_ptr, |
| 61 | # [num_reqs, num_speculative_steps + 1] |
| 62 | sampled_ptr, |
| 63 | sampled_stride, |
| 64 | # [num_reqs] |
| 65 | num_sampled_ptr, |
| 66 | # [num_reqs + 1] |
| 67 | cu_num_logits_ptr, |
| 68 | ): |
| 69 | req_idx = tl.program_id(0) |
| 70 | start_idx = tl.load(cu_num_logits_ptr + req_idx) |
| 71 | num_sampled = tl.load(num_sampled_ptr + req_idx) |
| 72 | for i in range(num_sampled): |
| 73 | token_id = tl.load(sampled_ptr + req_idx * sampled_stride + i) |
| 74 | tl.store(flat_sampled_ptr + start_idx + i, token_id) |
| 75 | |
| 76 | |
| 77 | @triton.jit |
| 78 | def _gather_draft_sampled_kernel( |
| 79 | # [num_logits] |
| 80 | draft_sampled_ptr, |
| 81 | # [num_logits] |
| 82 | pos_ptr, |
| 83 | # [num_tokens] |
| 84 | input_ids_ptr, |
| 85 | # [num_tokens] |
| 86 | positions_ptr, |
| 87 | # [num_logits] |
| 88 | logits_indices_ptr, |
| 89 | # [num_logits] |
| 90 | expanded_idx_mapping_ptr, |
| 91 | # [num_logits] |
| 92 | expanded_local_pos_ptr, |
| 93 | # [max_num_reqs] |
| 94 | prefill_len_ptr, |
| 95 | num_logits, |
| 96 | BLOCK_SIZE: tl.constexpr, |
| 97 | ): |
| 98 | block = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) |
| 99 | mask = block < num_logits |
| 100 | token_idx = tl.load(logits_indices_ptr + block, mask=mask) |
| 101 | token_id = tl.load(input_ids_ptr + token_idx, mask=mask) |
| 102 | pos = tl.load(positions_ptr + token_idx, mask=mask) |
| 103 | req_state_idx = tl.load(expanded_idx_mapping_ptr + block, mask=mask) |
| 104 | local_pos = tl.load(expanded_local_pos_ptr + block, mask=mask) |
| 105 | prefill_len = tl.load(prefill_len_ptr + req_state_idx, mask=mask) |
| 106 | # Drafts are only proposed after sampling past the prefill, so any draft |
| 107 | # slots in a step starting within the prefill are placeholders (e.g. the |
| 108 | # padded first step after a P/D remote KV load). |
| 109 | is_placeholder = (local_pos > 0) & (pos - local_pos < prefill_len) |
| 110 | token_id = tl.where(is_placeholder, -1, token_id) |
| 111 | tl.store(draft_sampled_ptr + block, token_id, mask=mask) |
| 112 | tl.store(pos_ptr + block, pos, mask=mask) |
| 113 | |
| 114 | |
| 115 | def gather_draft_sampled( |
| 116 | input_ids: torch.Tensor, |
| 117 | positions: torch.Tensor, |
| 118 | logits_indices: torch.Tensor, |
| 119 | expanded_idx_mapping: torch.Tensor, |
| 120 | expanded_local_pos: torch.Tensor, |
| 121 | prefill_len: torch.Tensor, |
| 122 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 123 | """Gather the input token and position of each logits row. |
| 124 | |
| 125 | Draft rows of requests that have not yet sampled past their prefill are |
| 126 | set to -1 so that the rejection kernels reject them. |
| 127 | """ |
| 128 | num_logits = logits_indices.shape[0] |
| 129 | draft_sampled = input_ids.new_empty(num_logits) |
| 130 | pos = positions.new_empty(num_logits) |
| 131 | BLOCK_SIZE = 1024 |
| 132 | _gather_draft_sampled_kernel[(triton.cdiv(num_logits, BLOCK_SIZE),)]( |
| 133 | draft_sampled, |
| 134 | pos, |
| 135 | input_ids, |
| 136 | positions, |
| 137 | logits_indices, |
| 138 | expanded_idx_mapping, |
| 139 | expanded_local_pos, |
| 140 | prefill_len, |
| 141 | num_logits, |
| 142 | BLOCK_SIZE=BLOCK_SIZE, |
| 143 | ) |
| 144 | return draft_sampled, pos |
| 145 | |
| 146 | |
| 147 | class RejectionSampler: |
| 148 | def __init__( |
| 149 | self, |
| 150 | sampler: Sampler, |
| 151 | spec_config: SpeculativeConfig, |
| 152 | device: torch.device, |
| 153 | *, |
| 154 | watermark_key: int | None = None, |
| 155 | ): |
| 156 | self.sampler = sampler |
| 157 | self.watermark_key = watermark_key |
| 158 | if watermark_key is not None: |
| 159 | assert isinstance(sampler, GPUWatermarkSampler) |
| 160 | self.num_speculative_steps = spec_config.num_speculative_tokens |
| 161 | self.enable_adaptive_verification = spec_config.enable_adaptive_verification |
| 162 | rejection_sample_method = spec_config.rejection_sample_method |
| 163 | self.use_block_verification: bool = False |
| 164 | self.synthetic_conditional_rates: torch.Tensor | None = None |
| 165 | if rejection_sample_method == "synthetic": |
| 166 | assert spec_config.synthetic_acceptance_rates is not None |
| 167 | self.synthetic_conditional_rates = torch.tensor( |
| 168 | unconditional_to_conditional_rates( |
| 169 | spec_config.synthetic_acceptance_rates |
| 170 | ), |
| 171 | dtype=torch.float32, |
| 172 | device=device, |
| 173 | ) |
| 174 | elif rejection_sample_method == "block": |
| 175 | self.use_block_verification = True |
| 176 | |
| 177 | def _get_logprobs_tensors( |
| 178 | self, |
| 179 | sampled: torch.Tensor, |
| 180 | num_sampled: torch.Tensor, |
| 181 | logits: torch.Tensor, |
| 182 | cu_num_logits: torch.Tensor, |
| 183 | cu_num_logits_np: np.ndarray, |
| 184 | max_num_logprobs: int, |
| 185 | ) -> LogprobsTensors | None: |
| 186 | if max_num_logprobs == NO_LOGPROBS: |
| 187 | return None |
| 188 | |
| 189 | num_reqs = cu_num_logits.shape[0] - 1 |
| 190 | num_logits = logits.shape[0] |
| 191 | flat_sampled = torch.zeros( |
| 192 | num_logits, dtype=sampled.dtype, device=sampled.device |
| 193 | ) |
| 194 | _flatten_sampled_kernel[(num_reqs,)]( |
| 195 | flat_sampled, |
| 196 | sampled, |
| 197 | sampled.stride(0), |
| 198 | num_sampled, |
| 199 | cu_num_logits, |
| 200 | num_warps=1, |
| 201 | ) |
| 202 | expanded_logits = num_logits != num_reqs |
| 203 | cu_num_generated_tokens: list[int] | torch.Tensor | None = None |
| 204 | if expanded_logits: |
| 205 | if self.enable_adaptive_verification: |
| 206 | # Adaptive verification keeps the true per-request boundaries |
| 207 | # on device only; cu_num_logits_np holds the pre-compacted layout. |
| 208 | cu_num_generated_tokens = cu_num_logits.clone() |
| 209 | else: |
| 210 | cu_num_generated_tokens = cu_num_logits_np.tolist() |
| 211 | return compute_topk_scores( |
| 212 | logits, |
| 213 | max_num_logprobs, |
| 214 | flat_sampled, |
| 215 | cu_num_generated_tokens, |
| 216 | logits_mode=self.sampler.logprobs_mode |
| 217 | in ("raw_logits", "processed_logits"), |
| 218 | ) |
| 219 | |
| 220 | def _watermarking_kwargs( |
| 221 | self, |
| 222 | draft_sampled: torch.Tensor, |
| 223 | expanded_idx_mapping: torch.Tensor, |
| 224 | expanded_local_pos: torch.Tensor, |
| 225 | ) -> dict[str, Any]: |
| 226 | if self.watermark_key is None: |
| 227 | return {} |
| 228 | assert isinstance(self.sampler, GPUWatermarkSampler) |
| 229 | contexts = self.sampler._get_contexts( |
| 230 | expanded_idx_mapping, expanded_local_pos, draft_sampled |
| 231 | ) |
| 232 | watermarking_skip_mask = None |
| 233 | if self.sampler.deduplicate_contexts != "none": |
| 234 | watermarking_skip_mask = self.sampler._get_repeated_contexts( |
| 235 | expanded_idx_mapping, contexts, expanded_local_pos |
| 236 | ) |
| 237 | return { |
| 238 | "contexts": contexts, |
| 239 | "watermarking": self.sampler.watermarking.gpu, |
| 240 | "watermarking_skip_mask": watermarking_skip_mask, |
| 241 | "watermark_key": self.watermark_key, |
| 242 | } |
| 243 | |
| 244 | def _verify( |
| 245 | self, |
| 246 | logits: torch.Tensor, |
| 247 | draft_logits: torch.Tensor | None, |
| 248 | draft_sampled: torch.Tensor, |
| 249 | pos: torch.Tensor, |
| 250 | cu_num_logits: torch.Tensor, |
| 251 | idx_mapping: torch.Tensor, |
| 252 | idx_mapping_np: np.ndarray, |
| 253 | expanded_idx_mapping: torch.Tensor, |
| 254 | expanded_local_pos: torch.Tensor, |
| 255 | seq_lens_upper_bound_np: np.ndarray, |
| 256 | verify_draft_sampled: torch.Tensor, |
| 257 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 258 | processed_logits = self.sampler.apply_sampling_params( |
| 259 | logits, |
| 260 | expanded_idx_mapping, |
| 261 | idx_mapping, |
| 262 | idx_mapping_np, |
| 263 | pos, |
| 264 | draft_sampled, |
| 265 | expanded_local_pos, |
| 266 | seq_lens_upper_bound_np, |
| 267 | ) |
| 268 | sampled, num_sampled = rejection_sample( |
| 269 | processed_logits, |
| 270 | draft_logits, |
| 271 | verify_draft_sampled, |
| 272 | cu_num_logits, |
| 273 | pos, |
| 274 | idx_mapping, |
| 275 | expanded_idx_mapping, |
| 276 | expanded_local_pos, |
| 277 | self.sampler.sampling_states.temperature.gpu, |
| 278 | self.sampler.sampling_states.seeds.gpu, |
| 279 | self.num_speculative_steps, |
| 280 | self.synthetic_conditional_rates, |
| 281 | use_fp64=self.sampler.use_fp64_gumbel, |
| 282 | use_block_verification=self.use_block_verification, |
| 283 | **self._watermarking_kwargs( |
| 284 | draft_sampled, expanded_idx_mapping, expanded_local_pos |
| 285 | ), |
| 286 | ) |
| 287 | return processed_logits, sampled, num_sampled |
| 288 | |
| 289 | def _verify_in_chunks( |
| 290 | self, |
| 291 | logits: torch.Tensor, |
| 292 | input_batch: InputBatch, |
| 293 | draft_logits: torch.Tensor | None, |
| 294 | draft_sampled: torch.Tensor, |
| 295 | pos: torch.Tensor, |
| 296 | max_chunk_logits: int, |
| 297 | max_num_logprobs: int, |
| 298 | verify_draft_sampled: torch.Tensor | None = None, |
| 299 | ) -> tuple[ |
| 300 | torch.Tensor, |
| 301 | torch.Tensor, |
| 302 | LogprobsTensors | None, |
| 303 | SamplingMaskTensors | None, |
| 304 | ]: |
| 305 | if verify_draft_sampled is None: |
| 306 | verify_draft_sampled = draft_sampled |
| 307 | cu_num_logits_np = input_batch.cu_num_logits_np |
| 308 | use_processed_logits = self.sampler.logprobs_mode in PROCESSED_LOGPROBS_MODES |
| 309 | num_reqs = input_batch.num_reqs |
| 310 | |
| 311 | if logits.shape[0] <= max_chunk_logits: |
| 312 | # One chunk covers the batch. Adaptive verification compacts the logits |
| 313 | # without updating cu_num_logits_np (it keeps the pre-compacted layout), |
| 314 | # so the stale sums must not pick chunk boundaries; its budget cap |
| 315 | # guarantees the compacted batch always lands here. |
| 316 | request_chunks: Iterable[tuple[int, int]] = ((0, num_reqs),) |
| 317 | else: |
| 318 | assert not self.enable_adaptive_verification |
| 319 | request_chunks = _iter_request_chunks(cu_num_logits_np, max_chunk_logits) |
| 320 | |
| 321 | sampled_chunks: list[torch.Tensor] = [] |
| 322 | num_sampled_chunks: list[torch.Tensor] = [] |
| 323 | logprobs_chunks: list[LogprobsTensors] = [] |
| 324 | sampling_mask_chunks: list[SamplingMaskTensors] = [] |
| 325 | max_num_kept = None |
| 326 | if self.sampler.return_sampling_mask: |
| 327 | max_num_kept = int( |
| 328 | np.max( |
| 329 | self.sampler.sampling_states.top_k.np[input_batch.idx_mapping_np] |
| 330 | ) |
| 331 | ) |
| 332 | |
| 333 | for start, end in request_chunks: |
| 334 | lo = int(cu_num_logits_np[start]) |
| 335 | hi = int(cu_num_logits_np[end]) |
| 336 | chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo |
| 337 | chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo |
| 338 | # draft_logits uses persistent request-state indices and stays global. |
| 339 | processed_logits, sampled, num_sampled = self._verify( |
| 340 | logits[lo:hi], |
| 341 | draft_logits, |
| 342 | draft_sampled[lo:hi], |
| 343 | pos[lo:hi], |
| 344 | chunk_cu_num_logits, |
| 345 | input_batch.idx_mapping[start:end], |
| 346 | input_batch.idx_mapping_np[start:end], |
| 347 | input_batch.expanded_idx_mapping[lo:hi], |
| 348 | input_batch.expanded_local_pos[lo:hi], |
| 349 | input_batch.seq_lens_cpu_upper_bound.numpy()[start:end], |
| 350 | verify_draft_sampled[lo:hi], |
| 351 | ) |
| 352 | chunk_logprobs = self._get_logprobs_tensors( |
| 353 | sampled, |
| 354 | num_sampled, |
| 355 | processed_logits if use_processed_logits else logits[lo:hi], |
| 356 | chunk_cu_num_logits, |
| 357 | chunk_cu_num_logits_np, |
| 358 | max_num_logprobs, |
| 359 | ) |
| 360 | if chunk_logprobs is not None: |
| 361 | logprobs_chunks.append(chunk_logprobs) |
| 362 | if max_num_kept is not None: |
| 363 | sampling_mask_chunks.append( |
| 364 | SamplingMaskTensors.from_logits( |
| 365 | processed_logits, |
| 366 | chunk_cu_num_logits, |
| 367 | num_sampled, |
| 368 | max_num_kept, |
| 369 | self.num_speculative_steps + 1, |
| 370 | ) |
| 371 | ) |
| 372 | del processed_logits |
| 373 | sampled_chunks.append(sampled) |
| 374 | num_sampled_chunks.append(num_sampled) |
| 375 | |
| 376 | if len(sampled_chunks) == 1: |
| 377 | logprobs_tensors = logprobs_chunks[0] if logprobs_chunks else None |
| 378 | sampling_mask_tensors = ( |
| 379 | sampling_mask_chunks[0] if sampling_mask_chunks else None |
| 380 | ) |
| 381 | return ( |
| 382 | sampled_chunks[0], |
| 383 | num_sampled_chunks[0], |
| 384 | logprobs_tensors, |
| 385 | sampling_mask_tensors, |
| 386 | ) |
| 387 | |
| 388 | logprobs_tensors = None |
| 389 | if logprobs_chunks: |
| 390 | expanded_logits = logits.shape[0] != input_batch.num_reqs |
| 391 | logprobs_tensors = LogprobsTensors.cat( |
| 392 | logprobs_chunks, |
| 393 | cu_num_generated_tokens=( |
| 394 | cu_num_logits_np.tolist() if expanded_logits else None |
| 395 | ), |
| 396 | ) |
| 397 | |
| 398 | sampled = torch.cat(sampled_chunks) |
| 399 | num_sampled = torch.cat(num_sampled_chunks) |
| 400 | sampling_mask_tensors = ( |
| 401 | SamplingMaskTensors.cat(sampling_mask_chunks) |
| 402 | if sampling_mask_chunks |
| 403 | else None |
| 404 | ) |
| 405 | return sampled, num_sampled, logprobs_tensors, sampling_mask_tensors |
| 406 | |
| 407 | def __call__( |
| 408 | self, |
| 409 | logits: torch.Tensor, |
| 410 | input_batch: InputBatch, |
| 411 | draft_logits: torch.Tensor | None = None, |
| 412 | invalid_drafts: torch.Tensor | None = None, |
| 413 | ) -> SamplerOutput: |
| 414 | # NOTE(woosuk): We intentionally compute num_nans before sampling to make clear |
| 415 | # that num_nans is computed before applying penalties and temperature. |
| 416 | num_nans = get_num_nans(logits) if self.sampler.compute_nans else None |
| 417 | |
| 418 | draft_sampled, pos = gather_draft_sampled( |
| 419 | input_batch.input_ids, |
| 420 | input_batch.positions, |
| 421 | input_batch.logits_indices, |
| 422 | input_batch.expanded_idx_mapping, |
| 423 | input_batch.expanded_local_pos, |
| 424 | self.sampler.req_states.prefill_len.gpu, |
| 425 | ) |
| 426 | |
| 427 | # Marking a draft invalid (`is_valid_draft = draft_sampled >= 0`) pins the |
| 428 | # accepted length before it reaches a permissive bitmask row. Only |
| 429 | # verification sees the copy: `apply_sampling_params` and watermarking |
| 430 | # keep the real draft ids. |
| 431 | verify_draft_sampled = draft_sampled |
| 432 | if invalid_drafts is not None: |
| 433 | verify_draft_sampled = draft_sampled.masked_fill(invalid_drafts, -1) |
| 434 | |
| 435 | max_num_logprobs = self.sampler.sampling_states.max_num_logprobs( |
| 436 | input_batch.idx_mapping_np |
| 437 | ) |
| 438 | chunk_logit_limit = get_max_chunk_logits(logits.shape[1]) |
| 439 | ( |
| 440 | sampled, |
| 441 | num_sampled, |
| 442 | logprobs_tensors, |
| 443 | sampling_mask_tensors, |
| 444 | ) = self._verify_in_chunks( |
| 445 | logits, |
| 446 | input_batch, |
| 447 | draft_logits, |
| 448 | draft_sampled, |
| 449 | pos, |
| 450 | chunk_logit_limit, |
| 451 | max_num_logprobs, |
| 452 | verify_draft_sampled, |
| 453 | ) |
| 454 | |
| 455 | num_sampled, num_rejected = get_num_sampled_and_rejected( |
| 456 | num_sampled, |
| 457 | input_batch.seq_lens, |
| 458 | input_batch.cu_num_logits, |
| 459 | input_batch.idx_mapping, |
| 460 | self.sampler.req_states.prefill_len.gpu, |
| 461 | ) |
| 462 | |
| 463 | return SamplerOutput( |
| 464 | sampled_token_ids=sampled, |
| 465 | logprobs_tensors=logprobs_tensors, |
| 466 | num_nans=num_nans, |
| 467 | num_sampled=num_sampled, |
| 468 | num_rejected=num_rejected, |
| 469 | sampling_mask_tensors=sampling_mask_tensors, |
| 470 | ) |