| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | |
| 5 | from vllm.config import VllmConfig |
| 6 | from vllm.v1.worker.gpu_input_batch import InputBatch |
| 7 | |
| 8 | |
| 9 | class SuffixDecodingProposer: |
| 10 | """Speculative decoding proposer for Suffix Decoding (https://arxiv.org/pdf/2411.04975). |
| 11 | This class imports and uses the official implementation from Arctic Inference |
| 12 | (https://github.com/snowflakedb/ArcticInference). |
| 13 | """ |
| 14 | |
| 15 | def __init__(self, vllm_config: VllmConfig): |
| 16 | config = vllm_config.speculative_config |
| 17 | assert config is not None, "Speculative config must be set" |
| 18 | self.num_speculative_tokens = config.num_speculative_tokens |
| 19 | self.max_tree_depth = config.suffix_decoding_max_tree_depth |
| 20 | self.max_spec_factor = config.suffix_decoding_max_spec_factor |
| 21 | self.min_token_prob = config.suffix_decoding_min_token_prob |
| 22 | self.max_model_len = vllm_config.model_config.max_model_len |
| 23 | |
| 24 | # Lazy import to avoid error when Suffix Decoding is not used. |
| 25 | from arctic_inference.suffix_decoding import SuffixDecodingCache |
| 26 | |
| 27 | # Initialize and empty cache. This object will take care of caching request |
| 28 | # outputs, evicting old requests, and manages the per-prompt suffix trees. |
| 29 | self.suffix_cache = SuffixDecodingCache( |
| 30 | max_tree_depth=config.suffix_decoding_max_tree_depth, |
| 31 | max_cached_requests=config.suffix_decoding_max_cached_requests, |
| 32 | ) |
| 33 | |
| 34 | def propose( |
| 35 | self, |
| 36 | num_speculative_tokens: int, |
| 37 | input_batch: InputBatch, |
| 38 | sampled_token_ids: list[list[int]], |
| 39 | slot_mappings: dict[str, torch.Tensor] |
| 40 | | list[dict[str, torch.Tensor]] |
| 41 | | None = None, # unused |
| 42 | ) -> list[list[int]]: |
| 43 | assert num_speculative_tokens == self.num_speculative_tokens |
| 44 | """ |
| 45 | Propose speculative tokens for each request in the input batch. Suffix Decoding |
| 46 | will speculate a dynamic number of tokens for each request every decoding step, |
| 47 | so each entry in the returned list may have different lengths. |
| 48 | """ |
| 49 | draft_token_ids: list[list[int]] = [] |
| 50 | for i, sampled_ids in enumerate(sampled_token_ids): |
| 51 | if not sampled_ids: |
| 52 | # Skip speculative decoding for partial prefills. |
| 53 | draft_token_ids.append([]) |
| 54 | continue |
| 55 | |
| 56 | req_id = input_batch.req_ids[i] |
| 57 | num_tokens = input_batch.num_tokens_no_spec[i] |
| 58 | if num_tokens >= self.max_model_len: |
| 59 | # Skip requests that have already reached the max model length. |
| 60 | draft_token_ids.append([]) |
| 61 | continue |
| 62 | |
| 63 | index = input_batch.req_id_to_index[req_id] |
| 64 | if req_id not in self.suffix_cache.active_requests: |
| 65 | if req_id in self.suffix_cache.cached_requests: |
| 66 | # Reset the suffix cache for this request. |
| 67 | self.suffix_cache.evict_cached_response(req_id) |
| 68 | num_prompt_tokens = input_batch.num_prompt_tokens[index] |
| 69 | prompt_token_ids = input_batch.token_ids_cpu[index, :num_prompt_tokens] |
| 70 | # Start a new request, this will build the suffix tree for that prompt. |
| 71 | self.suffix_cache.start_request(req_id, prompt_token_ids) |
| 72 | |
| 73 | # Append the newly sampled ids to the suffix cache for this request. |
| 74 | self.suffix_cache.add_active_response(req_id, sampled_ids) |
| 75 | |
| 76 | # Suffix decoding only uses the most recent tokens up to max_tree_depth, so |
| 77 | # we extract the pattern from the end of the input. |
| 78 | start = max(0, num_tokens - self.max_tree_depth) |
| 79 | pattern = input_batch.token_ids_cpu[i, start:num_tokens] |
| 80 | draft = self.suffix_cache.speculate( |
| 81 | req_id, |
| 82 | pattern, |
| 83 | max_spec_tokens=min( |
| 84 | self.num_speculative_tokens, self.max_model_len - num_tokens - 1 |
| 85 | ), |
| 86 | max_spec_factor=self.max_spec_factor, |
| 87 | min_token_prob=self.min_token_prob, |
| 88 | ) |
| 89 | |
| 90 | draft_token_ids.append(draft.token_ids) |
| 91 | |
| 92 | # Stop requests that were not seen in the input batch. |
| 93 | for req_id in ( |
| 94 | self.suffix_cache.active_requests - input_batch.req_id_to_index.keys() |
| 95 | ): |
| 96 | self.suffix_cache.stop_request(req_id) |
| 97 | |
| 98 | return draft_token_ids |
| 99 | |
| 100 | def load_model(self, *args, **kwargs): |
| 101 | # No model to load. |
| 102 | pass |