| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from abc import ABC, abstractmethod |
| 4 | from collections.abc import Mapping |
| 5 | from dataclasses import replace |
| 6 | from typing import TYPE_CHECKING, Any |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | import torch.nn as nn |
| 11 | |
| 12 | from vllm.config import VllmConfig, get_layers_from_vllm_config |
| 13 | from vllm.config.compilation import CUDAGraphMode |
| 14 | from vllm.distributed.eplb.eplb_state import EplbState |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 17 | from vllm.model_executor.models import supports_multimodal_embeddings |
| 18 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 19 | from vllm.v1.watermarking import create_watermarker |
| 20 | from vllm.v1.watermarking.spec_decode import ( |
| 21 | DraftWatermarker, |
| 22 | create_speculative_draft_watermarker, |
| 23 | ) |
| 24 | from vllm.v1.worker.gpu.attn_utils import ( |
| 25 | build_attn_metadata, |
| 26 | init_attn_backend, |
| 27 | ) |
| 28 | from vllm.v1.worker.gpu.block_table import BlockTables |
| 29 | from vllm.v1.worker.gpu.cp_utils import prepare_dcp_local_seq_lens |
| 30 | from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor |
| 31 | from vllm.v1.worker.gpu.dp_utils import DPSyncState |
| 32 | from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers |
| 33 | from vllm.v1.worker.gpu.model_states.interface import ModelState |
| 34 | from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample |
| 35 | from vllm.v1.worker.gpu.spec_decode.acceptance_estimator import ( |
| 36 | OnlineAcceptanceEstimator, |
| 37 | ) |
| 38 | from vllm.v1.worker.utils import AttentionGroup |
| 39 | |
| 40 | if TYPE_CHECKING: |
| 41 | from vllm.v1.watermarking.gpu_sampler import GPUWatermarkSampler |
| 42 | from vllm.v1.worker.gpu.pcp_manager import PCPManager |
| 43 | |
| 44 | logger = init_logger(__name__) |
| 45 | |
| 46 | |
| 47 | def _target_feeds_hc_residual(vllm_config: VllmConfig) -> bool: |
| 48 | """Whether the target replaces the drafter's input with its HC residual. |
| 49 | |
| 50 | Keyed on the same hook the model runner calls to perform that swap. It is |
| 51 | resolved from the target model class because speculators are built before |
| 52 | the target model is instantiated. |
| 53 | """ |
| 54 | from vllm.model_executor.model_loader.utils import get_model_cls |
| 55 | |
| 56 | target_cls = get_model_cls(vllm_config.model_config) |
| 57 | return hasattr(target_cls, "get_mtp_target_hidden_states") |
| 58 | |
| 59 | |
| 60 | class BaseSpeculator(ABC): |
| 61 | num_query_per_req: int = 1 |
| 62 | # Extra query slots reserved per request outside the regular queries. |
| 63 | num_extra_query_per_req: int = 0 |
| 64 | |
| 65 | @abstractmethod |
| 66 | def init_cudagraph_manager(self, cudagraph_mode: CUDAGraphMode) -> None: |
| 67 | pass |
| 68 | |
| 69 | @abstractmethod |
| 70 | def capture(self) -> None: |
| 71 | pass |
| 72 | |
| 73 | @abstractmethod |
| 74 | def propose( |
| 75 | self, |
| 76 | input_batch: InputBatch, |
| 77 | attn_metadata: dict[str, Any], |
| 78 | slot_mappings: dict[str, torch.Tensor], |
| 79 | # [num_tokens, hidden_size] |
| 80 | last_hidden_states: torch.Tensor, |
| 81 | # num_layers x [num_tokens, hidden_size] |
| 82 | aux_hidden_states: list[torch.Tensor] | None, |
| 83 | # [num_reqs] |
| 84 | num_sampled: torch.Tensor, |
| 85 | # [num_reqs] |
| 86 | num_rejected: torch.Tensor, |
| 87 | # [max_num_reqs] |
| 88 | last_sampled: torch.Tensor, |
| 89 | # [num_prefill_lookahead, max_num_reqs] |
| 90 | next_prefill_tokens: torch.Tensor, |
| 91 | # [max_num_reqs] |
| 92 | temperature: torch.Tensor, |
| 93 | # [max_num_reqs] |
| 94 | seeds: torch.Tensor, |
| 95 | dp_sync: DPSyncState | None = None, |
| 96 | dummy_run: bool = False, |
| 97 | skip_attn_for_dummy_run: bool = False, |
| 98 | mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, |
| 99 | is_profile: bool = False, |
| 100 | ) -> torch.Tensor: |
| 101 | pass |
| 102 | |
| 103 | |
| 104 | class DraftModelSpeculator(BaseSpeculator): |
| 105 | def __init__(self, vllm_config: VllmConfig, device: torch.device): |
| 106 | self.vllm_config = vllm_config |
| 107 | self.device = device |
| 108 | |
| 109 | assert vllm_config.speculative_config is not None |
| 110 | self.speculative_config = vllm_config.speculative_config |
| 111 | self.method = self.speculative_config.method |
| 112 | self.num_speculative_steps = self.speculative_config.num_speculative_tokens |
| 113 | self.draft_model_config = self.speculative_config.draft_model_config |
| 114 | |
| 115 | self.scheduler_config = vllm_config.scheduler_config |
| 116 | self.max_num_reqs = self.scheduler_config.max_num_seqs |
| 117 | self.max_num_tokens = self.scheduler_config.max_num_batched_tokens |
| 118 | self.max_model_len = vllm_config.model_config.max_model_len |
| 119 | self.draft_max_seq_len = self.max_model_len |
| 120 | # We need to get the hidden size from the draft model config because |
| 121 | # the draft model's hidden size can be different from the target model's |
| 122 | # hidden size (e.g., Llama 3.3 70B). |
| 123 | self.hidden_size = self.draft_model_config.get_hidden_size() |
| 124 | # Widen for HC-multiplexed residuals: a target that implements |
| 125 | # get_mtp_target_hidden_states() (e.g. DeepSeek V4) hands the drafter |
| 126 | # its pre-hc_head (T, hc_mult * hidden_size) residual in place of the |
| 127 | # collapsed hidden states, so the drafter's buffers must match. Key off |
| 128 | # that hook rather than hc_mult alone -- HY V4 runs iHC in its backbone |
| 129 | # (hc_mult=4) but its MTP head consumes the collapsed states, so |
| 130 | # widening it feeds propose() a 4x-too-wide buffer. |
| 131 | if _target_feeds_hc_residual(vllm_config): |
| 132 | hc_mult = getattr(self.draft_model_config.hf_config, "hc_mult", 1) |
| 133 | self.hidden_size = self.hidden_size * hc_mult |
| 134 | self.vocab_size = self.draft_model_config.get_vocab_size() |
| 135 | self.dtype = vllm_config.model_config.dtype |
| 136 | self.use_fp64_gumbel = vllm_config.model_config.use_fp64_gumbel |
| 137 | self.use_local_argmax_reduction = ( |
| 138 | self.speculative_config.use_local_argmax_reduction |
| 139 | ) |
| 140 | |
| 141 | # DP configuration |
| 142 | self.dp_size = vllm_config.parallel_config.data_parallel_size |
| 143 | self.dp_rank = vllm_config.parallel_config.data_parallel_rank |
| 144 | self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size |
| 145 | |
| 146 | self.eplb_state: EplbState | None = None |
| 147 | |
| 148 | self.input_buffers = InputBuffers( |
| 149 | max_num_reqs=self.max_num_reqs, |
| 150 | max_num_tokens=self.max_num_tokens, |
| 151 | device=device, |
| 152 | ) |
| 153 | self.idx_mapping = torch.zeros( |
| 154 | self.max_num_reqs, dtype=torch.int32, device=device |
| 155 | ) |
| 156 | self.temperature = torch.zeros( |
| 157 | self.max_num_reqs, dtype=torch.float32, device=device |
| 158 | ) |
| 159 | self.seeds = torch.zeros(self.max_num_reqs, dtype=torch.int64, device=device) |
| 160 | self.draft_tokens = torch.zeros( |
| 161 | self.max_num_reqs, |
| 162 | self.num_speculative_steps, |
| 163 | dtype=torch.int64, |
| 164 | device=device, |
| 165 | ) |
| 166 | self.arange_np = np.arange(self.max_num_reqs + 1, dtype=np.int32) |
| 167 | self.draft_is_prefilling = torch.zeros(self.max_num_reqs, dtype=torch.bool) |
| 168 | |
| 169 | self.draft_token_confidence_probs = torch.empty_like( |
| 170 | self.draft_tokens, dtype=torch.float32 |
| 171 | ) |
| 172 | self.enable_adaptive_verification = ( |
| 173 | self.speculative_config.enable_adaptive_verification |
| 174 | ) |
| 175 | self.use_acceptance_estimator = self.enable_adaptive_verification |
| 176 | self.acceptance_estimator: OnlineAcceptanceEstimator | None = None |
| 177 | |
| 178 | self.draft_logits: torch.Tensor | None = None |
| 179 | if self.speculative_config.draft_sample_method == "probabilistic": |
| 180 | # Pre-temperature logits, cached from the previous decode step. |
| 181 | dtype, fill = self.draft_logits_spec(vllm_config) |
| 182 | self.draft_logits = torch.full( |
| 183 | ( |
| 184 | self.max_num_reqs, |
| 185 | self.num_speculative_steps, |
| 186 | self.vocab_size, |
| 187 | ), |
| 188 | fill, |
| 189 | dtype=dtype, |
| 190 | device=device, |
| 191 | ) |
| 192 | |
| 193 | self.draft_watermarker: DraftWatermarker | None = None |
| 194 | watermark_config = getattr(vllm_config, "watermark_config", None) |
| 195 | if watermark_config is not None: |
| 196 | watermarker = create_watermarker(watermark_config) |
| 197 | self.draft_watermarker = create_speculative_draft_watermarker( |
| 198 | watermarker, |
| 199 | self.max_num_reqs, |
| 200 | device, |
| 201 | watermark_config.allow_target_only_watermarking, |
| 202 | self.num_speculative_steps, |
| 203 | watermark_config.deduplicate_contexts, |
| 204 | watermark_config.deduplicate_contexts_max_history, |
| 205 | ) |
| 206 | |
| 207 | self.supports_mm_inputs = False |
| 208 | self.pcp_manager: PCPManager | None = None |
| 209 | |
| 210 | @abstractmethod |
| 211 | def load_draft_model( |
| 212 | self, target_model: nn.Module, target_attn_layer_names: set[str] |
| 213 | ) -> nn.Module: |
| 214 | pass |
| 215 | |
| 216 | def load_model(self, target_model: nn.Module) -> None: |
| 217 | target_attn_layer_names = set( |
| 218 | get_layers_from_vllm_config( |
| 219 | self.vllm_config, |
| 220 | AttentionLayerBase, # type: ignore[type-abstract] |
| 221 | ).keys() |
| 222 | ) |
| 223 | |
| 224 | self.model = self.load_draft_model(target_model, target_attn_layer_names) |
| 225 | self._validate_local_argmax_reduction() |
| 226 | |
| 227 | all_attn_layers = get_layers_from_vllm_config( |
| 228 | self.vllm_config, |
| 229 | AttentionLayerBase, # type: ignore[type-abstract] |
| 230 | ) |
| 231 | self.draft_attn_layer_names = set(all_attn_layers) - target_attn_layer_names |
| 232 | for layer_name in self.draft_attn_layer_names: |
| 233 | all_attn_layers[layer_name].is_draft_layer = True |
| 234 | |
| 235 | target_supports_mm = self.vllm_config.model_config.supports_multimodal_inputs |
| 236 | draft_supports_mm = supports_multimodal_embeddings(self.model) |
| 237 | self.supports_mm_inputs = target_supports_mm and draft_supports_mm |
| 238 | if target_supports_mm and not draft_supports_mm: |
| 239 | logger.warning_once( |
| 240 | "Draft model %s does not support external multimodal embeddings. " |
| 241 | "Embeddings from the target model will not be passed to the " |
| 242 | "drafter; using text-only draft inputs instead.", |
| 243 | type(self.model).__name__, |
| 244 | ) |
| 245 | |
| 246 | if self.use_acceptance_estimator: |
| 247 | if self.use_local_argmax_reduction: |
| 248 | raise ValueError( |
| 249 | "Adaptive verification without a confidence head estimates " |
| 250 | "per-position acceptance from the draft logits, which " |
| 251 | "use_local_argmax_reduction never materializes. Disable one " |
| 252 | "of them." |
| 253 | ) |
| 254 | self.acceptance_estimator = OnlineAcceptanceEstimator( |
| 255 | self.max_num_reqs, |
| 256 | self.num_speculative_steps, |
| 257 | self.device, |
| 258 | ) |
| 259 | |
| 260 | def set_eplb_state(self, eplb_state: EplbState) -> None: |
| 261 | """Inject EPLB state after construction.""" |
| 262 | self.eplb_state = eplb_state |
| 263 | |
| 264 | def _prepare_eplb_forward(self, num_unpadded_tokens: int) -> None: |
| 265 | """Call EPLB prepare_forward if EPLB is active for the draft model.""" |
| 266 | if self.eplb_state is not None: |
| 267 | self.eplb_state.prepare_forward( |
| 268 | self.speculative_config.draft_model_config, |
| 269 | num_unpadded_tokens, |
| 270 | ) |
| 271 | |
| 272 | @property |
| 273 | def attn_vllm_config(self) -> VllmConfig: |
| 274 | """Config for the draft's attention metadata builders. Overridden by |
| 275 | speculators whose attention mode differs from the target's.""" |
| 276 | return self.vllm_config |
| 277 | |
| 278 | def set_attn( |
| 279 | self, |
| 280 | model_state: ModelState, |
| 281 | kv_cache_config: KVCacheConfig, |
| 282 | block_tables: BlockTables, |
| 283 | target_input_buffers: InputBuffers, |
| 284 | target_attn_groups: list[list[AttentionGroup]], |
| 285 | ) -> None: |
| 286 | self.model_state = model_state |
| 287 | self.kv_cache_config = kv_cache_config |
| 288 | self.attn_groups, self.attn_cg_support, _ = init_attn_backend( |
| 289 | kv_cache_config, |
| 290 | self.attn_vllm_config, |
| 291 | self.device, |
| 292 | active_layer_names=self.draft_attn_layer_names, |
| 293 | ) |
| 294 | self.block_tables = block_tables |
| 295 | # The target model runner's buffers and attention groups. Draft |
| 296 | # prefill reuses the target model's attention metadata, so its |
| 297 | # cudagraph capture must build dummy metadata through the same |
| 298 | # builders and buffers. |
| 299 | self.target_input_buffers = target_input_buffers |
| 300 | self.target_attn_groups = target_attn_groups |
| 301 | |
| 302 | def _build_attn_metadata( |
| 303 | self, |
| 304 | num_reqs: int, |
| 305 | batch_desc: BatchExecutionDescriptor, |
| 306 | query_start_loc_np: np.ndarray, |
| 307 | seq_lens_cpu_upper_bound: torch.Tensor, |
| 308 | step: int, |
| 309 | causal: bool | Mapping[int, bool] = True, |
| 310 | dcp_local_seq_lens: torch.Tensor | None = None, |
| 311 | slot_mappings: torch.Tensor | None = None, |
| 312 | ) -> dict[str, Any] | None: |
| 313 | num_reqs_padded = batch_desc.num_reqs or num_reqs |
| 314 | # A FULL graph replays a captured shape whose padded requests each hold |
| 315 | # a full query width, so attention must see the padded token count. |
| 316 | # PIECEWISE/eager needs the actual token count, because batch_desc may |
| 317 | # carry graph or DP padding that no request owns, which would desync it |
| 318 | # from query_start_loc. |
| 319 | num_tokens = ( |
| 320 | batch_desc.num_tokens |
| 321 | if batch_desc.cg_mode == CUDAGraphMode.FULL |
| 322 | else int(query_start_loc_np[-1]) |
| 323 | ) |
| 324 | query_start_loc_cpu = torch.empty(num_reqs_padded + 1, dtype=torch.int32) |
| 325 | query_start_loc_cpu[: num_reqs + 1] = torch.from_numpy( |
| 326 | query_start_loc_np[: num_reqs + 1] |
| 327 | ) |
| 328 | query_start_loc_cpu[num_reqs:] = query_start_loc_cpu[num_reqs] |
| 329 | max_query_len = int((query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).max()) |
| 330 | block_tables = [ |
| 331 | x[:num_reqs_padded] for x in self.block_tables.input_block_tables |
| 332 | ] |
| 333 | if slot_mappings is None: |
| 334 | slot_mappings = self.block_tables.slot_mappings[:, :num_tokens] |
| 335 | draft_seq_lens_cpu_upper_bound = torch.zeros( |
| 336 | num_reqs_padded, dtype=torch.int32, device="cpu" |
| 337 | ) |
| 338 | torch.add( |
| 339 | seq_lens_cpu_upper_bound[:num_reqs], |
| 340 | step, |
| 341 | out=draft_seq_lens_cpu_upper_bound[:num_reqs], |
| 342 | ) |
| 343 | draft_seq_lens_cpu_upper_bound[:num_reqs].clamp_(max=self.max_model_len) |
| 344 | if dcp_local_seq_lens is None and self.dcp_size > 1: |
| 345 | # Draft steps advance and rewind their own global sequence lengths, |
| 346 | # so the target model's DCP-local lengths may already be stale. |
| 347 | dcp_local_seq_lens = prepare_dcp_local_seq_lens( |
| 348 | self.input_buffers.dcp_local_seq_lens, |
| 349 | self.input_buffers.seq_lens, |
| 350 | num_reqs, |
| 351 | self.dcp_size, |
| 352 | self.block_tables.cp_rank, |
| 353 | self.block_tables.cp_interleave, |
| 354 | ) |
| 355 | attn_metadata = build_attn_metadata( |
| 356 | attn_groups=self.attn_groups, |
| 357 | num_reqs=num_reqs_padded, |
| 358 | num_tokens=num_tokens, |
| 359 | query_start_loc_gpu=self.input_buffers.query_start_loc[ |
| 360 | : num_reqs_padded + 1 |
| 361 | ], |
| 362 | query_start_loc_cpu=query_start_loc_cpu, |
| 363 | max_query_len=max_query_len, |
| 364 | seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], |
| 365 | dcp_local_seq_lens=( |
| 366 | None |
| 367 | if dcp_local_seq_lens is None |
| 368 | else dcp_local_seq_lens[:num_reqs_padded] |
| 369 | ), |
| 370 | max_seq_len=self.draft_max_seq_len, |
| 371 | block_tables=block_tables, |
| 372 | slot_mappings=slot_mappings, |
| 373 | kv_cache_config=self.kv_cache_config, |
| 374 | causal=causal, |
| 375 | seq_lens_cpu_upper_bound=draft_seq_lens_cpu_upper_bound, |
| 376 | is_prefilling=self.draft_is_prefilling[:num_reqs_padded], |
| 377 | ) |
| 378 | return attn_metadata |
| 379 | |
| 380 | def draft_logits_spec(self, vllm_config: VllmConfig) -> tuple[torch.dtype, float]: |
| 381 | """Dtype and fill for the cached proposal distribution. |
| 382 | |
| 383 | Speculators that write only a subset of columns each step override this. |
| 384 | """ |
| 385 | return vllm_config.model_config.head_dtype, 0.0 |
| 386 | |
| 387 | def _validate_local_argmax_reduction(self) -> None: |
| 388 | if not self.use_local_argmax_reduction: |
| 389 | return |
| 390 | if self.speculative_config.draft_sample_method == "probabilistic": |
| 391 | raise ValueError( |
| 392 | "use_local_argmax_reduction is not compatible with " |
| 393 | "draft_sample_method='probabilistic'." |
| 394 | ) |
| 395 | if not hasattr(self.model, "get_top_tokens"): |
| 396 | raise ValueError( |
| 397 | "use_local_argmax_reduction is enabled but draft model " |
| 398 | f"{self.model.__class__.__name__} does not implement " |
| 399 | "get_top_tokens()." |
| 400 | ) |
| 401 | logger.info( |
| 402 | "Using local argmax reduction for draft token generation " |
| 403 | "(communication: O(2*tp_size) vs O(vocab_size))." |
| 404 | ) |
| 405 | |
| 406 | def compute_draft_logits( |
| 407 | self, hidden_states: torch.Tensor, spec_step_idx: int |
| 408 | ) -> torch.Tensor: |
| 409 | return self.model.compute_logits(hidden_states) |
| 410 | |
| 411 | def get_draft_top_tokens( |
| 412 | self, hidden_states: torch.Tensor, spec_step_idx: int |
| 413 | ) -> torch.Tensor: |
| 414 | return self.model.get_top_tokens(hidden_states) |
| 415 | |
| 416 | def sample_draft( |
| 417 | self, |
| 418 | hidden_states: torch.Tensor, |
| 419 | sample_src_positions: torch.Tensor, |
| 420 | idx_mapping: torch.Tensor, |
| 421 | temperature: torch.Tensor, |
| 422 | seeds: torch.Tensor, |
| 423 | draft_step: torch.Tensor, |
| 424 | draft_logits: torch.Tensor | None, |
| 425 | spec_step_idx: int = 0, |
| 426 | ) -> torch.Tensor: |
| 427 | if draft_logits is None and self.use_local_argmax_reduction: |
| 428 | return self.get_draft_top_tokens(hidden_states, spec_step_idx) |
| 429 | |
| 430 | logits = self.compute_draft_logits(hidden_states, spec_step_idx) |
| 431 | if draft_logits is not None: |
| 432 | sampler = ( |
| 433 | gumbel_sample |
| 434 | if self.draft_watermarker is None |
| 435 | else self.draft_watermarker.sample |
| 436 | ) |
| 437 | sampled = sampler( |
| 438 | logits, |
| 439 | idx_mapping, |
| 440 | temperature, |
| 441 | seeds, |
| 442 | sample_src_positions, |
| 443 | apply_temperature=True, |
| 444 | is_drafting=True, |
| 445 | logits_cache=draft_logits, |
| 446 | logits_cache_col=draft_step, |
| 447 | use_fp64=self.use_fp64_gumbel, |
| 448 | ) |
| 449 | else: |
| 450 | sampled = logits.argmax(dim=-1) |
| 451 | self._maybe_predict_acceptance(logits, idx_mapping, draft_step) |
| 452 | return sampled |
| 453 | |
| 454 | def _maybe_predict_acceptance( |
| 455 | self, logits: torch.Tensor, idx_mapping: torch.Tensor, draft_step: torch.Tensor |
| 456 | ) -> None: |
| 457 | if self.acceptance_estimator is not None: |
| 458 | self.acceptance_estimator.predict( |
| 459 | logits, |
| 460 | idx_mapping, |
| 461 | draft_step, |
| 462 | self.draft_token_confidence_probs, |
| 463 | self.temperature, |
| 464 | ) |
| 465 | |
| 466 | def observe_verification( |
| 467 | self, |
| 468 | idx_mapping: torch.Tensor, |
| 469 | num_sampled: torch.Tensor, |
| 470 | num_rejected: torch.Tensor, |
| 471 | ) -> None: |
| 472 | """Fold the target's verdict on the last drafts into the estimator. |
| 473 | |
| 474 | Must run before the next `propose`, which overwrites the per-slot |
| 475 | features the verdict grades. |
| 476 | """ |
| 477 | if self.acceptance_estimator is not None: |
| 478 | self.acceptance_estimator.step(idx_mapping, num_sampled, num_rejected) |
| 479 | |
| 480 | def prepare_watermarking( |
| 481 | self, |
| 482 | sampler: "GPUWatermarkSampler", |
| 483 | idx_mapping: torch.Tensor, |
| 484 | ) -> None: |
| 485 | if self.draft_watermarker is None: |
| 486 | return |
| 487 | self.draft_watermarker.prepare( |
| 488 | sampler._get_contexts(idx_mapping), |
| 489 | sampler.watermarking.gpu[idx_mapping], |
| 490 | sampler.req_states.all_token_ids.gpu, |
| 491 | sampler.req_states.prompt_len.gpu, |
| 492 | sampler.req_states.total_len.gpu, |
| 493 | ) |
| 494 | |
| 495 | def _copy_request_inputs( |
| 496 | self, |
| 497 | num_reqs: int, |
| 498 | # [num_reqs] |
| 499 | idx_mapping: torch.Tensor, |
| 500 | # [max_num_reqs] |
| 501 | temperature: torch.Tensor, |
| 502 | # [max_num_reqs] |
| 503 | seeds: torch.Tensor, |
| 504 | dummy_run: bool = False, |
| 505 | ) -> None: |
| 506 | # Copy temperature, seeds, and idx mapping to the pre-allocated buffers. |
| 507 | # NOTE(woosuk): For draft sampling, we only consider the temperature |
| 508 | # and ignore the other sampling parameters such as top_k and top_p, |
| 509 | # for simplicity and performance. |
| 510 | # While this may slightly degrade the acceptance rate, it does not |
| 511 | # affect the output distribution after rejection sampling. |
| 512 | self.temperature.copy_(temperature) |
| 513 | self.seeds.copy_(seeds) |
| 514 | # idx_mapping == -1 marks a row the drafter must not act on: sampling |
| 515 | # skips it and compute_slot_mappings emits PAD. CUDA-graph padded rows |
| 516 | # always get it; a dummy batch gets it for every row, since its arange |
| 517 | # idx_mapping names request-state slots with stale block tables. |
| 518 | if dummy_run: |
| 519 | self.idx_mapping.fill_(-1) |
| 520 | else: |
| 521 | self.idx_mapping[:num_reqs].copy_(idx_mapping) |
| 522 | self.idx_mapping[num_reqs:].fill_(-1) |
| 523 | |
| 524 | def _build_uniform_batch_dp_sync( |
| 525 | self, |
| 526 | target_dp_sync: DPSyncState, |
| 527 | num_reqs: int, |
| 528 | num_query_per_req: int = 1, |
| 529 | ) -> tuple[DPSyncState, int]: |
| 530 | num_batch_tokens = target_dp_sync.num_reqs * num_query_per_req |
| 531 | assert num_reqs * num_query_per_req <= num_batch_tokens, ( |
| 532 | "reusing a DP sync that does not cover this batch's requests" |
| 533 | ) |
| 534 | return replace( |
| 535 | target_dp_sync, |
| 536 | num_tokens_across_dp=torch.full_like( |
| 537 | target_dp_sync.num_tokens_across_dp, num_batch_tokens |
| 538 | ), |
| 539 | uniform_token_count=num_query_per_req, |
| 540 | eager=False, |
| 541 | ), num_batch_tokens |
| 542 | |
| 543 | def _build_uniform_attn_metadata( |
| 544 | self, |
| 545 | batch_desc: BatchExecutionDescriptor, |
| 546 | num_reqs: int, |
| 547 | num_query_per_req: int, |
| 548 | seq_lens_cpu_upper_bound: torch.Tensor, |
| 549 | step: int, |
| 550 | causal: bool | Mapping[int, bool] = True, |
| 551 | dcp_local_seq_lens: torch.Tensor | None = None, |
| 552 | ) -> dict[str, Any] | None: |
| 553 | query_start_loc_np = self.arange_np[: num_reqs + 1] * num_query_per_req |
| 554 | return self._build_attn_metadata( |
| 555 | num_reqs=num_reqs, |
| 556 | batch_desc=batch_desc, |
| 557 | query_start_loc_np=query_start_loc_np, |
| 558 | seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, |
| 559 | step=step, |
| 560 | causal=causal, |
| 561 | dcp_local_seq_lens=dcp_local_seq_lens, |
| 562 | ) |