| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import numpy as np |
| 4 | import torch |
| 5 | |
| 6 | from vllm.config import LoadConfig, replace |
| 7 | from vllm.distributed import get_pp_group |
| 8 | from vllm.logger import init_logger |
| 9 | from vllm.v1.outputs import DraftTokenIds |
| 10 | from vllm.v1.worker.gpu.async_utils import async_copy_to_np |
| 11 | from vllm.v1.worker.gpu.input_batch import InputBatch |
| 12 | |
| 13 | logger = init_logger(__name__) |
| 14 | |
| 15 | |
| 16 | def get_pp_safe_draft_load_config(load_config: LoadConfig) -> LoadConfig: |
| 17 | """Avoid collectives that include PP ranks without a draft model.""" |
| 18 | if get_pp_group().world_size > 1 and load_config.load_format == "fastsafetensors": |
| 19 | logger.warning_once( |
| 20 | "fastsafetensors cannot load a draft model instantiated on only " |
| 21 | "one pipeline stage; falling back to the standard safetensors " |
| 22 | "loader for the draft model. The target model still uses " |
| 23 | "fastsafetensors." |
| 24 | ) |
| 25 | return replace(load_config, load_format="auto") |
| 26 | return load_config |
| 27 | |
| 28 | |
| 29 | class DraftTokensHandler: |
| 30 | def __init__(self, device: torch.device | None = None): |
| 31 | self.device = device |
| 32 | self.copy_stream = torch.cuda.Stream(device) |
| 33 | # Blocking (sleep) event to avoid busy-polling the CUDA driver lock. |
| 34 | self.copy_event = torch.cuda.Event(blocking=True) |
| 35 | |
| 36 | self.req_ids: list[str] = [] |
| 37 | self.draft_tokens_np: np.ndarray | None = None |
| 38 | self.num_draft_tokens: int = 0 |
| 39 | |
| 40 | def set_draft_tokens( |
| 41 | self, input_batch: InputBatch, draft_tokens: torch.Tensor |
| 42 | ) -> None: |
| 43 | self.req_ids = input_batch.req_ids |
| 44 | self.num_draft_tokens = draft_tokens.shape[1] |
| 45 | if not input_batch.has_structured_output_reqs: |
| 46 | # No draft token validation needs to be performed by |
| 47 | # the scheduler for this batch. |
| 48 | self.draft_tokens_np = None |
| 49 | return |
| 50 | |
| 51 | # For spec decoding + structured outputs, we must transfer the |
| 52 | # draft tokens back to the scheduler for grammar validation. |
| 53 | current_stream = torch.cuda.current_stream(self.device) |
| 54 | self.copy_stream.wait_stream(current_stream) |
| 55 | with torch.cuda.stream(self.copy_stream): |
| 56 | self.draft_tokens_np = async_copy_to_np(draft_tokens) |
| 57 | # draft_tokens is a temporary allocation on the main stream and read here on |
| 58 | # copy_stream; without record_stream, the caching allocator may reuse its |
| 59 | # memory before the async copy executes. |
| 60 | draft_tokens.record_stream(self.copy_stream) |
| 61 | self.copy_event.record() |
| 62 | |
| 63 | def get_draft_tokens(self) -> DraftTokenIds | None: |
| 64 | if self.draft_tokens_np is not None: |
| 65 | self.copy_event.synchronize() |
| 66 | draft_token_ids = self.draft_tokens_np.tolist() |
| 67 | else: |
| 68 | # This case only happens when async scheduling is disabled. |
| 69 | draft_token_ids = [[-1] * self.num_draft_tokens for _ in self.req_ids] |
| 70 | return DraftTokenIds(self.req_ids, draft_token_ids) |
| 71 | |
| 72 | |
| 73 | def get_parallel_drafting_token_id(hf_config) -> int: |
| 74 | """Resolve the mask token id used for parallel drafting slots. |
| 75 | |
| 76 | Checks (in order): `dflash_config.mask_token_id`, top-level `mask_token_id`, |
| 77 | `dspark_noise_token_id`, `pard_token`, `ptd_token_id`. Raises ValueError if |
| 78 | none are present. |
| 79 | """ |
| 80 | dflash_config = getattr(hf_config, "dflash_config", None) or {} |
| 81 | if "mask_token_id" in dflash_config: |
| 82 | return int(dflash_config["mask_token_id"]) |
| 83 | if getattr(hf_config, "mask_token_id", None) is not None: |
| 84 | return int(hf_config.mask_token_id) |
| 85 | if hasattr(hf_config, "dspark_noise_token_id"): |
| 86 | return int(hf_config.dspark_noise_token_id) |
| 87 | if hasattr(hf_config, "pard_token"): |
| 88 | return int(hf_config.pard_token) |
| 89 | if hasattr(hf_config, "ptd_token_id"): |
| 90 | return int(hf_config.ptd_token_id) |
| 91 | raise ValueError( |
| 92 | "Model config must specify `dflash_config.mask_token_id`," |
| 93 | " `mask_token_id`, `dspark_noise_token_id`, `pard_token`, or" |
| 94 | " `ptd_token_id` for parallel drafting." |
| 95 | ) |