| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Adaptive verification for DSpark speculative decoding.""" |
| 4 | |
| 5 | from collections import defaultdict |
| 6 | from collections.abc import Iterable, Iterator |
| 7 | from typing import TYPE_CHECKING |
| 8 | |
| 9 | import numpy as np |
| 10 | import torch |
| 11 | |
| 12 | import vllm.envs as envs |
| 13 | from vllm.config.compilation import CUDAGraphMode |
| 14 | from vllm.distributed.parallel_state import get_tp_group |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.utils.gpu_sync_debug import gpu_sync_allowed |
| 17 | from vllm.utils.torch_utils import async_tensor_h2d |
| 18 | from vllm.v1.attention.backend import AttentionCGSupport |
| 19 | from vllm.v1.utils import CpuGpuBuffer |
| 20 | from vllm.v1.worker.gpu.async_utils import StepTimingSample, stream |
| 21 | from vllm.v1.worker.gpu.attn_utils import ( |
| 22 | get_query_lens_mismatch_unsupported_backend, |
| 23 | get_varlen_cudagraph_unsupported_backend, |
| 24 | ) |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | _PROFILE_REPLAYS = 5 |
| 28 | |
| 29 | if TYPE_CHECKING: |
| 30 | from vllm.config import VllmConfig |
| 31 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 32 | from vllm.v1.worker.gpu.states import RequestState |
| 33 | from vllm.v1.worker.utils import AttentionGroup |
| 34 | |
| 35 | |
| 36 | def resolve_adaptive_cudagraph_mode( |
| 37 | mode: CUDAGraphMode, *, piecewise_capture_available: bool |
| 38 | ) -> CUDAGraphMode: |
| 39 | """Select a separate decode route for variable-length verification.""" |
| 40 | assert mode.has_full_cudagraphs() |
| 41 | if mode == CUDAGraphMode.FULL_DECODE_ONLY or not piecewise_capture_available: |
| 42 | return CUDAGraphMode.FULL_DECODE_ONLY |
| 43 | return CUDAGraphMode.FULL_AND_PIECEWISE |
| 44 | |
| 45 | |
| 46 | def _assign_draft_token_budget( |
| 47 | confidence_probs: torch.Tensor, |
| 48 | idx_mapping: torch.Tensor, |
| 49 | capacities: torch.Tensor, |
| 50 | draft_budget: int, |
| 51 | num_steps: int, |
| 52 | ) -> None: |
| 53 | """Admit the globally best `draft_budget` draft slots, in place. |
| 54 | |
| 55 | Every (request, step) slot is scored by its survival probability, the running |
| 56 | product of that request's per-position confidences, and the highest scores win. |
| 57 | Survival only decreases along a request, so a global top-k always admits |
| 58 | continuously along steps with a request. |
| 59 | |
| 60 | `capacities` enters holding each request's scheduled draft count (which bounds its |
| 61 | eligible slots) and leaves holding the admitted count. The caller only calls this |
| 62 | when `draft_budget < capacities.sum()`, so every winner is a real draft slot. |
| 63 | """ |
| 64 | survival = confidence_probs[idx_mapping].cumprod(dim=1) |
| 65 | steps = torch.arange(num_steps, device=survival.device) |
| 66 | # Out-of-range slots score -inf so they never outrank a real draft. |
| 67 | out_of_range = steps[None, :] >= capacities[:, None] |
| 68 | survival = survival.masked_fill(out_of_range, -float("inf")) |
| 69 | flat = survival.flatten() |
| 70 | winners = flat.topk(draft_budget).indices |
| 71 | admitted = torch.zeros_like(flat, dtype=torch.bool).index_fill_(0, winners, True) |
| 72 | torch.sum(admitted.view_as(survival), dim=1, dtype=capacities.dtype, out=capacities) |
| 73 | |
| 74 | |
| 75 | _assign_draft_token_budget_compiled = torch.compile( |
| 76 | _assign_draft_token_budget, dynamic=True |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | def build_cost_tables_from_curves( |
| 81 | draft_curve: list[tuple[int, float]], |
| 82 | verify_curve: list[tuple[int, float]], |
| 83 | max_num_reqs: int, |
| 84 | max_batch_tokens: int, |
| 85 | cudagraph_limit: int = 0, |
| 86 | ) -> tuple[np.ndarray, np.ndarray]: |
| 87 | """Build cost tables: graph-padded below the capture limit, smooth above. |
| 88 | |
| 89 | Args: |
| 90 | draft_curve: (size, cost) samples for the draft model. |
| 91 | verify_curve: (size, cost) samples for the verify model. |
| 92 | max_num_reqs: Largest request count to build a table for. |
| 93 | max_batch_tokens: Largest token count to build a table for. |
| 94 | cudagraph_limit: Largest cudagraph-captured size. At or below it, |
| 95 | execution pads up to the next captured size, so cost is a step |
| 96 | function. Above it there is no padding, so cost is continuous. |
| 97 | |
| 98 | """ |
| 99 | |
| 100 | def build_table(limit: int, curve: list[tuple[int, float]]) -> np.ndarray: |
| 101 | xs, ys = np.asarray(curve, dtype=np.float64).T |
| 102 | ys = np.maximum.accumulate(ys) |
| 103 | values = np.arange(limit + 1) |
| 104 | # Execution pads to the next captured size, so cost is a step |
| 105 | # function of the padded size: smooth interpolation would invent |
| 106 | # marginal per-token costs that don't exist within a pad bucket. |
| 107 | idx = np.searchsorted(xs, values, side="left") |
| 108 | result = ys[np.minimum(idx, len(xs) - 1)] |
| 109 | # Past the capture limit nothing pads, so cost really is continuous in |
| 110 | # size; snapping to the next profiled point overestimates badly when |
| 111 | # the profiled points are far apart. Interpolate only between points |
| 112 | # that are themselves past the limit: crossing the limit loses |
| 113 | # cudagraphs entirely, which is a genuine discontinuity, so the first |
| 114 | # point above it must not be blended with the last one below. |
| 115 | if cudagraph_limit: |
| 116 | smooth = values > cudagraph_limit |
| 117 | above = xs > cudagraph_limit |
| 118 | if smooth.any() and above.any(): |
| 119 | result[smooth] = np.interp(values[smooth], xs[above], ys[above]) |
| 120 | if len(xs) > 1: |
| 121 | after = values > xs[-1] |
| 122 | slope = (ys[-1] - ys[-2]) / (xs[-1] - xs[-2]) |
| 123 | result[after] = ys[-1] + slope * (values[after] - xs[-1]) |
| 124 | return result |
| 125 | |
| 126 | draft_table = np.maximum(build_table(max_num_reqs, draft_curve), 0.0) |
| 127 | verify_table = np.maximum(build_table(max_batch_tokens, verify_curve), 1e-6) |
| 128 | return draft_table, verify_table |
| 129 | |
| 130 | |
| 131 | class AdaptiveVerificationManager: |
| 132 | def __init__( |
| 133 | self, |
| 134 | req_states: "RequestState", |
| 135 | query_start_loc: torch.Tensor, |
| 136 | num_bonus_tokens: int, |
| 137 | max_total_logits: int, |
| 138 | ): |
| 139 | self.req_states = req_states |
| 140 | self.num_speculative_steps = req_states.num_speculative_steps |
| 141 | device = req_states.device |
| 142 | self._copy_stream = torch.cuda.Stream(device) |
| 143 | |
| 144 | self.num_bonus_tokens = num_bonus_tokens |
| 145 | # Rejection sampling verifies logits in one contiguous chunk; the |
| 146 | # chunked path indexes by scheduled (untrimmed) offsets and cannot |
| 147 | # address the compacted layout, so the budget must fit one chunk. |
| 148 | self._max_total_logits = max_total_logits |
| 149 | self.query_start_loc = query_start_loc |
| 150 | self.cost_tables: tuple[np.ndarray, np.ndarray] | None = None |
| 151 | # Largest cudagraph-captured token count; above it nothing pads. |
| 152 | self._cudagraph_limit = 0 |
| 153 | self._batch_budget: tuple[dict[str, int], dict[str, int], int] | None = None |
| 154 | max_num_reqs = req_states.max_num_reqs |
| 155 | # Current per-slot confidences |
| 156 | self._confidence_probs = torch.empty( |
| 157 | (max_num_reqs, self.num_speculative_steps), |
| 158 | dtype=torch.float32, |
| 159 | device=device, |
| 160 | ) |
| 161 | self._batch_draft_capacity = torch.empty( |
| 162 | max_num_reqs, dtype=torch.int32, device=device |
| 163 | ) |
| 164 | self._num_non_draft_tokens = torch.empty_like(query_start_loc[:-1]) |
| 165 | self._cu_num_logits = torch.empty_like(query_start_loc) |
| 166 | |
| 167 | # Two D2H slots preserve stale inputs for budget selection. |
| 168 | self._stale_confidences = [ |
| 169 | CpuGpuBuffer( |
| 170 | max_num_reqs, |
| 171 | self.num_speculative_steps, |
| 172 | dtype=torch.float32, |
| 173 | device=device, |
| 174 | ) |
| 175 | for _ in range(2) |
| 176 | ] |
| 177 | self._copy_events = [torch.cuda.Event(blocking=True) for _ in range(2)] |
| 178 | self._pending_resets: list[int] = [] |
| 179 | self._stale_idx = 0 |
| 180 | for slot in self._stale_confidences: |
| 181 | slot.np.fill(1.0) |
| 182 | |
| 183 | def add_request(self, req_idx: int) -> None: |
| 184 | self._stale_confidences[self._stale_idx].np[req_idx].fill(1.0) |
| 185 | self._pending_resets.append(req_idx) |
| 186 | self._confidence_probs[req_idx].fill_(1.0) |
| 187 | |
| 188 | def batches_to_profile(self, capture_sizes: list[int]) -> Iterator[dict[str, int]]: |
| 189 | """Dummy-run kwargs whose step timings seed the cost tables. |
| 190 | |
| 191 | Run these inside StepTimingCollector.collect(), then hand the block's |
| 192 | timings to set_initial_cost_curves.""" |
| 193 | max_num_tokens = self.req_states.max_num_batched_tokens |
| 194 | size = self._cudagraph_limit = capture_sizes[-1] if capture_sizes else 0 |
| 195 | # Also profile beyond the capture limit: real steps run there |
| 196 | # (piecewise/eager) and linear extrapolation badly underestimates |
| 197 | # them. These runs double as JIT warmup for the piecewise shapes. |
| 198 | tail_sizes: set[int] = set() |
| 199 | if size: |
| 200 | tail_sizes.add(min(size + size // 2, max_num_tokens)) |
| 201 | while size < max_num_tokens: |
| 202 | size = min(size * 2, max_num_tokens) |
| 203 | tail_sizes.add(size) |
| 204 | tail_sizes -= set(capture_sizes) |
| 205 | for num_tokens in capture_sizes + sorted(tail_sizes): |
| 206 | for _ in range(_PROFILE_REPLAYS): |
| 207 | yield { |
| 208 | "num_tokens": num_tokens, |
| 209 | "context_len": envs.VLLM_ADAPTIVE_VERIFICATION_PROFILE_CONTEXT_LEN, |
| 210 | } |
| 211 | |
| 212 | def set_initial_cost_curves(self, samples: list[StepTimingSample]) -> None: |
| 213 | def median_curve( |
| 214 | points: Iterable[tuple[int, float]], |
| 215 | ) -> list[tuple[int, float]]: |
| 216 | grouped: defaultdict[int, list[float]] = defaultdict(list) |
| 217 | for key, value in points: |
| 218 | grouped[key].append(value) |
| 219 | return [(k, float(np.median(v))) for k, v in sorted(grouped.items())] |
| 220 | |
| 221 | # Draft curve: eager-target steps inflate drafter timings (the CPU is |
| 222 | # still launching kernels, opening gaps between the drafter's events), |
| 223 | # and request counts — unlike token counts — collide across execution |
| 224 | # modes, so only graph-replay samples may price the draft curve. |
| 225 | draft_curve = median_curve( |
| 226 | (s.num_reqs, s.drafter_ms) for s in samples if s.full_cudagraph |
| 227 | ) |
| 228 | verify_curve = median_curve( |
| 229 | (s.num_target_tokens, s.forward_ms) for s in samples |
| 230 | ) |
| 231 | self.set_cost_curves(draft_curve, verify_curve) |
| 232 | |
| 233 | def set_cost_curves( |
| 234 | self, |
| 235 | draft_curve: list[tuple[int, float]], |
| 236 | verify_curve: list[tuple[int, float]], |
| 237 | ) -> None: |
| 238 | draft_curve, verify_curve = get_tp_group().broadcast_object( |
| 239 | (draft_curve, verify_curve), src=0 |
| 240 | ) |
| 241 | if not draft_curve or not verify_curve: |
| 242 | raise RuntimeError( |
| 243 | "Adaptive verification could not profile step costs. Pass " |
| 244 | "`enable_adaptive_verification=false` in the speculative config to " |
| 245 | "verify a fixed number of drafts instead." |
| 246 | ) |
| 247 | self.cost_tables = build_cost_tables_from_curves( |
| 248 | draft_curve, |
| 249 | verify_curve, |
| 250 | self.req_states.max_num_reqs, |
| 251 | self.req_states.max_num_batched_tokens, |
| 252 | self._cudagraph_limit, |
| 253 | ) |
| 254 | logger.debug("DSpark cost tables: %s", self.cost_tables) |
| 255 | |
| 256 | def record_confidences( |
| 257 | self, confidence_probs: torch.Tensor, input_batch: "InputBatch" |
| 258 | ) -> None: |
| 259 | """Publish this step's raw confidences for the ranking kernel and start |
| 260 | copying them to the CPU, where a later step's budget reads them.""" |
| 261 | num_reqs = input_batch.num_reqs |
| 262 | ready_idx = self._stale_idx ^ 1 |
| 263 | with gpu_sync_allowed(): |
| 264 | self._copy_events[ready_idx].synchronize() |
| 265 | if self._pending_resets: |
| 266 | self._stale_confidences[ready_idx].np[self._pending_resets] = 1.0 |
| 267 | self._pending_resets.clear() |
| 268 | # Last step's copy has landed: budgets read it, this step overwrites the |
| 269 | # slot they were reading before. |
| 270 | self._stale_idx, write_idx = ready_idx, self._stale_idx |
| 271 | |
| 272 | self._confidence_probs[input_batch.idx_mapping] = confidence_probs[:num_reqs] |
| 273 | write_slot = self._stale_confidences[write_idx] |
| 274 | write_slot.gpu.copy_(self._confidence_probs) |
| 275 | |
| 276 | current_stream = torch.cuda.current_stream(self.req_states.device) |
| 277 | self._copy_stream.wait_stream(current_stream) |
| 278 | with stream(self._copy_stream, current_stream): |
| 279 | write_slot.copy_to_cpu() |
| 280 | self._copy_events[write_idx].record() |
| 281 | |
| 282 | def get_num_tokens( |
| 283 | self, |
| 284 | num_tokens_per_req: dict[str, int], |
| 285 | draft_tokens: dict[str, list[int]], |
| 286 | ) -> int: |
| 287 | """Token count once the draft budget is trimmed to fit. |
| 288 | |
| 289 | Stashes the chosen budget in ``_batch_budget`` for the compaction and |
| 290 | reallocation that follow in the same step. |
| 291 | """ |
| 292 | assert self.cost_tables is not None |
| 293 | req_ids = list(num_tokens_per_req) |
| 294 | num_reqs = len(req_ids) |
| 295 | scheduled_tokens = np.fromiter( |
| 296 | num_tokens_per_req.values(), dtype=np.int32, count=num_reqs |
| 297 | ) |
| 298 | scheduled_drafts = np.fromiter( |
| 299 | (len(draft_tokens.get(req_id, ())) for req_id in req_ids), |
| 300 | dtype=np.int32, |
| 301 | count=num_reqs, |
| 302 | ) |
| 303 | num_non_draft_tokens = scheduled_tokens - scheduled_drafts |
| 304 | slots = np.fromiter( |
| 305 | map(self.req_states.req_id_to_index.__getitem__, req_ids), |
| 306 | dtype=np.int32, |
| 307 | count=num_reqs, |
| 308 | ) |
| 309 | stale_confidences = self._stale_confidences[self._stale_idx].np[slots] |
| 310 | survival_probability = np.cumprod(stale_confidences.astype(np.float64), axis=1) |
| 311 | steps = np.arange(self.num_speculative_steps) |
| 312 | valid = steps[None, :] < scheduled_drafts[:, None] |
| 313 | scores = np.sort(survival_probability[valid])[::-1] |
| 314 | num_non_draft_tokens_total = int(num_non_draft_tokens.sum()) |
| 315 | max_draft_budget = min( |
| 316 | int(scheduled_drafts.sum()), |
| 317 | max(0, self._max_total_logits - num_reqs * self.num_bonus_tokens), |
| 318 | ) |
| 319 | scores = scores[:max_draft_budget] |
| 320 | draft_cost_ms, verify_cost_ms = self.cost_tables |
| 321 | num_sampling_requests = np.count_nonzero( |
| 322 | self.req_states.num_computed_tokens_np[slots] + num_non_draft_tokens |
| 323 | >= self.req_states.prefill_len.np[slots] |
| 324 | ) |
| 325 | num_tokens_to_estimated_accepted_tokens = np.concatenate( |
| 326 | ([num_sampling_requests], num_sampling_requests + np.cumsum(scores)) |
| 327 | ) |
| 328 | costs = ( |
| 329 | draft_cost_ms[len(req_ids)] |
| 330 | + verify_cost_ms[ |
| 331 | num_non_draft_tokens_total : num_non_draft_tokens_total |
| 332 | + max_draft_budget |
| 333 | + 1 |
| 334 | ] |
| 335 | ) |
| 336 | num_drafts_per_req = { |
| 337 | req_id: int(num_drafts) |
| 338 | for req_id, num_drafts in zip(req_ids, scheduled_drafts, strict=True) |
| 339 | } |
| 340 | num_non_draft_tokens_per_req = { |
| 341 | req_id: int(num_tokens) |
| 342 | for req_id, num_tokens in zip(req_ids, num_non_draft_tokens, strict=True) |
| 343 | } |
| 344 | draft_budget = int(np.argmax(num_tokens_to_estimated_accepted_tokens / costs)) |
| 345 | self._batch_budget = ( |
| 346 | num_drafts_per_req, |
| 347 | num_non_draft_tokens_per_req, |
| 348 | draft_budget, |
| 349 | ) |
| 350 | return sum(num_non_draft_tokens_per_req.values()) + draft_budget |
| 351 | |
| 352 | def compact_batch( |
| 353 | self, |
| 354 | num_draft_tokens_per_req: np.ndarray, |
| 355 | num_scheduled_tokens: np.ndarray, |
| 356 | cu_num_logits_np: np.ndarray, |
| 357 | ) -> tuple[np.ndarray, np.ndarray]: |
| 358 | """Compact the CPU batch to the chosen draft budget. |
| 359 | |
| 360 | Returns the compacted per-request token counts and the CPU cu_num_logits_np. |
| 361 | If the draft budget is 0, we can know cu_num_logits_np exactly, otherwise |
| 362 | its unchanged/an-upper-bound. |
| 363 | """ |
| 364 | batch_budget = self._batch_budget |
| 365 | assert batch_budget is not None |
| 366 | _, _, draft_budget = batch_budget |
| 367 | num_drafts = int(num_draft_tokens_per_req.sum()) |
| 368 | if draft_budget == num_drafts: |
| 369 | return num_scheduled_tokens, cu_num_logits_np |
| 370 | |
| 371 | num_non_draft_tokens = num_scheduled_tokens - num_draft_tokens_per_req |
| 372 | if draft_budget == 0: |
| 373 | # The draft budget is 0, so we can know cu_num_logits_np exactly. This helps |
| 374 | # when we would exceed the sampler logit chunk size. |
| 375 | num_reqs = num_scheduled_tokens.shape[0] |
| 376 | cu_num_logits_np = ( |
| 377 | np.arange(num_reqs + 1, dtype=cu_num_logits_np.dtype) |
| 378 | * self.num_bonus_tokens |
| 379 | ) |
| 380 | return num_non_draft_tokens, cu_num_logits_np |
| 381 | |
| 382 | is_verification_request = num_draft_tokens_per_req > 0 |
| 383 | num_verification_reqs = int(is_verification_request.sum()) |
| 384 | # sort_batch_req_ids keeps verification requests at the front. |
| 385 | assert np.all(is_verification_request[:num_verification_reqs]) |
| 386 | # for the CPU side buffer we distribute draft tokens evenly |
| 387 | draft_lens_cpu = np.zeros_like(num_non_draft_tokens) |
| 388 | draft_lens_cpu[:num_verification_reqs] = draft_budget // num_verification_reqs |
| 389 | draft_lens_cpu[: draft_budget % num_verification_reqs] += 1 |
| 390 | return num_non_draft_tokens + draft_lens_cpu, cu_num_logits_np |
| 391 | |
| 392 | def reallocate_drafts( |
| 393 | self, req_ids: list[str], idx_mapping: torch.Tensor |
| 394 | ) -> tuple[torch.Tensor, torch.Tensor, int]: |
| 395 | batch_budget, self._batch_budget = self._batch_budget, None |
| 396 | assert batch_budget is not None |
| 397 | num_drafts_per_req, num_non_draft_tokens_per_req, draft_budget = batch_budget |
| 398 | num_reqs = idx_mapping.shape[0] |
| 399 | scheduled_drafts = np.fromiter( |
| 400 | (num_drafts_per_req[req_id] for req_id in req_ids), |
| 401 | dtype=np.int32, |
| 402 | count=num_reqs, |
| 403 | ) |
| 404 | num_non_draft_tokens = np.fromiter( |
| 405 | (num_non_draft_tokens_per_req[req_id] for req_id in req_ids), |
| 406 | dtype=np.int32, |
| 407 | count=num_reqs, |
| 408 | ) |
| 409 | num_tokens = int(num_non_draft_tokens.sum()) + draft_budget |
| 410 | |
| 411 | # Rank draft slots by survival probability and admit the best prefix. |
| 412 | # capacities enters holding each request's valid draft count (the kernel |
| 413 | # uses it to bound eligible slots) and leaves holding the admitted count. |
| 414 | capacities = self._batch_draft_capacity[:num_reqs] |
| 415 | if draft_budget == 0: |
| 416 | capacities.zero_() |
| 417 | else: |
| 418 | async_tensor_h2d(scheduled_drafts, out=capacities) |
| 419 | if draft_budget < int(scheduled_drafts.sum()): |
| 420 | _assign_draft_token_budget_compiled( |
| 421 | self._confidence_probs, |
| 422 | idx_mapping, |
| 423 | capacities, |
| 424 | draft_budget, |
| 425 | self.num_speculative_steps, |
| 426 | ) |
| 427 | |
| 428 | num_non_draft_tokens_gpu = self._num_non_draft_tokens[:num_reqs] |
| 429 | async_tensor_h2d(num_non_draft_tokens, out=num_non_draft_tokens_gpu) |
| 430 | self._cu_num_logits[:1].zero_() |
| 431 | torch.cumsum( |
| 432 | capacities + self.num_bonus_tokens, |
| 433 | dim=0, |
| 434 | out=self._cu_num_logits[1 : num_reqs + 1], |
| 435 | ) |
| 436 | self.query_start_loc[:1].zero_() |
| 437 | torch.cumsum( |
| 438 | capacities + num_non_draft_tokens_gpu, |
| 439 | dim=0, |
| 440 | out=self.query_start_loc[1 : num_reqs + 1], |
| 441 | ) |
| 442 | self.query_start_loc[num_reqs + 1 :].fill_(num_tokens) |
| 443 | return ( |
| 444 | self._cu_num_logits[: num_reqs + 1], |
| 445 | self.query_start_loc, |
| 446 | draft_budget, |
| 447 | ) |
| 448 | |
| 449 | |
| 450 | def maybe_create_adaptive_verification_manager( |
| 451 | *, |
| 452 | enable_adaptive_verification: bool, |
| 453 | attn_groups: list[list["AttentionGroup"]], |
| 454 | req_states: "RequestState", |
| 455 | query_start_loc: torch.Tensor, |
| 456 | num_bonus_tokens: int, |
| 457 | max_total_logits: int, |
| 458 | vllm_config: "VllmConfig", |
| 459 | target_layer_names: set[str] | None = None, |
| 460 | additional_attn_cg_support: tuple[AttentionCGSupport, str | None] | None = None, |
| 461 | ) -> AdaptiveVerificationManager | None: |
| 462 | if not enable_adaptive_verification: |
| 463 | return None |
| 464 | |
| 465 | # The selector rejects unsupported backends, but models that |
| 466 | # hard-wire theirs (e.g. DeepSeek-V4) never go through it. |
| 467 | backend = get_query_lens_mismatch_unsupported_backend( |
| 468 | attn_groups, checked_layer_names=target_layer_names |
| 469 | ) |
| 470 | if backend is not None: |
| 471 | raise ValueError( |
| 472 | "Adaptive verification trims verification requests on device, which" |
| 473 | f" the {backend} attention backend does not support. Pass " |
| 474 | "enable_adaptive_verification=false in the speculative config, or " |
| 475 | "use a backend that does." |
| 476 | ) |
| 477 | |
| 478 | # The runner's decode_query_len, the width varlen decode graphs capture. |
| 479 | max_query_len = req_states.num_speculative_steps + num_bonus_tokens |
| 480 | unsupported: tuple[str | None, int | None] | None = ( |
| 481 | get_varlen_cudagraph_unsupported_backend( |
| 482 | attn_groups, |
| 483 | vllm_config, |
| 484 | max_query_len, |
| 485 | checked_layer_names=target_layer_names, |
| 486 | ) |
| 487 | ) |
| 488 | if unsupported is None and additional_attn_cg_support is not None: |
| 489 | # Groups built outside init_attn_backend report only their support |
| 490 | # level, and without a bound only ALWAYS replays varlen batches. |
| 491 | additional_support, additional_backend = additional_attn_cg_support |
| 492 | if additional_support != AttentionCGSupport.ALWAYS: |
| 493 | unsupported = additional_backend, None |
| 494 | if unsupported is not None: |
| 495 | backend, bound = unsupported |
| 496 | allowed = "none" if bound is None else f"at most {bound}" |
| 497 | raise ValueError( |
| 498 | "Adaptive verification replays decode cudagraphs whose per-request " |
| 499 | f"query lengths vary up to {max_query_len}, but {backend} allows " |
| 500 | f"{allowed}. Pass enable_adaptive_verification=false in the " |
| 501 | "speculative config, or use a backend that does." |
| 502 | ) |
| 503 | |
| 504 | return AdaptiveVerificationManager( |
| 505 | req_states, |
| 506 | query_start_loc, |
| 507 | num_bonus_tokens, |
| 508 | max_total_logits=max_total_logits, |
| 509 | ) |