| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import TYPE_CHECKING |
| 4 | |
| 5 | import torch |
| 6 | |
| 7 | import vllm.envs as envs |
| 8 | from vllm.sampling_params import SamplingParams |
| 9 | from vllm.triton_utils import tl, triton |
| 10 | from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor |
| 11 | from vllm.v1.worker.gpu.sample.logits_processor.interface import ( |
| 12 | LogitsContext, |
| 13 | LogitsProcessor, |
| 14 | LogitsProcRequestState, |
| 15 | ) |
| 16 | |
| 17 | if TYPE_CHECKING: |
| 18 | from vllm.config import VllmConfig |
| 19 | |
| 20 | |
| 21 | class BadWordsState(LogitsProcessor): |
| 22 | def __init__(self, vllm_config: "VllmConfig", req_states: LogitsProcRequestState): |
| 23 | self.req_states = req_states |
| 24 | max_num_reqs = req_states.max_num_reqs |
| 25 | device = req_states.device |
| 26 | |
| 27 | max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS |
| 28 | max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS |
| 29 | # flattened bad word tokens: [max_num_reqs, VLLM_MAX_BAD_WORDS_TOTAL_TOKENS] |
| 30 | self.bad_word_token_ids = StagedWriteTensor( |
| 31 | (max_num_reqs, max_total_tokens), dtype=torch.int32, device=device |
| 32 | ) |
| 33 | # cumulative offsets of bad words: [max_num_reqs, VLLM_MAX_NUM_BAD_WORDS + 1] |
| 34 | self.bad_word_offsets = StagedWriteTensor( |
| 35 | (max_num_reqs, max_num_bad_words + 1), dtype=torch.int32, device=device |
| 36 | ) |
| 37 | # number of bad words per request |
| 38 | self.num_bad_words = UvaBackedTensor(max_num_reqs, dtype=torch.int32) |
| 39 | |
| 40 | def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool: |
| 41 | bad_words_token_ids = sampling_params.bad_words_token_ids |
| 42 | if not bad_words_token_ids: |
| 43 | self.num_bad_words.np[req_idx] = 0 |
| 44 | return False |
| 45 | |
| 46 | num_bad_words = len(bad_words_token_ids) |
| 47 | max_num_bad_words = envs.VLLM_MAX_NUM_BAD_WORDS |
| 48 | if num_bad_words > max_num_bad_words: |
| 49 | raise ValueError( |
| 50 | f"Too many bad words: {num_bad_words}. " |
| 51 | f"The max number is {max_num_bad_words}." |
| 52 | ) |
| 53 | |
| 54 | # Flatten bad words and compute offsets |
| 55 | flattened_tokens: list[int] = [] |
| 56 | offsets: list[int] = [0] |
| 57 | for bad_word in bad_words_token_ids: |
| 58 | flattened_tokens.extend(bad_word) |
| 59 | offsets.append(len(flattened_tokens)) |
| 60 | |
| 61 | max_total_tokens = envs.VLLM_MAX_BAD_WORDS_TOTAL_TOKENS |
| 62 | if len(flattened_tokens) > max_total_tokens: |
| 63 | raise ValueError( |
| 64 | f"Too many total bad word tokens: {len(flattened_tokens)}. " |
| 65 | f"The max is {max_total_tokens}." |
| 66 | ) |
| 67 | |
| 68 | # Stage writes |
| 69 | self.bad_word_token_ids.stage_write(req_idx, 0, flattened_tokens) |
| 70 | self.bad_word_offsets.stage_write(req_idx, 0, offsets) |
| 71 | self.num_bad_words.np[req_idx] = num_bad_words |
| 72 | return True |
| 73 | |
| 74 | def apply_staged_writes(self) -> None: |
| 75 | self.num_bad_words.copy_to_uva() |
| 76 | self.bad_word_token_ids.apply_write() |
| 77 | self.bad_word_offsets.apply_write() |
| 78 | |
| 79 | def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor: |
| 80 | max_num_bad_words = int(self.num_bad_words.np[ctx.idx_mapping_np].max()) |
| 81 | if max_num_bad_words == 0: |
| 82 | # No request uses bad words. Skip the kernel launch. |
| 83 | return logits |
| 84 | |
| 85 | apply_bad_words( |
| 86 | logits, |
| 87 | ctx.expanded_idx_mapping, |
| 88 | self.bad_word_token_ids.gpu, |
| 89 | self.bad_word_offsets.gpu, |
| 90 | self.num_bad_words.gpu, |
| 91 | self.req_states.all_token_ids.gpu, |
| 92 | self.req_states.prompt_len.gpu, |
| 93 | self.req_states.total_len.gpu, |
| 94 | ctx.input_ids, |
| 95 | ctx.expanded_local_pos, |
| 96 | max_num_bad_words, |
| 97 | ) |
| 98 | return logits |
| 99 | |
| 100 | |
| 101 | @triton.jit |
| 102 | def _bad_words_kernel( |
| 103 | logits_ptr, |
| 104 | logits_stride, |
| 105 | expanded_idx_mapping_ptr, |
| 106 | bad_word_token_ids_ptr, |
| 107 | bad_word_token_ids_stride, |
| 108 | bad_word_offsets_ptr, |
| 109 | bad_word_offsets_stride, |
| 110 | num_bad_words_ptr, |
| 111 | all_token_ids_ptr, |
| 112 | all_token_ids_stride, |
| 113 | prompt_len_ptr, |
| 114 | total_len_ptr, |
| 115 | input_ids_ptr, |
| 116 | expanded_local_pos_ptr, |
| 117 | ): |
| 118 | token_idx = tl.program_id(0).to(tl.int64) |
| 119 | bw_idx = tl.program_id(1) |
| 120 | |
| 121 | req_state_idx = tl.load(expanded_idx_mapping_ptr + token_idx) |
| 122 | num_bad_words = tl.load(num_bad_words_ptr + req_state_idx) |
| 123 | |
| 124 | if bw_idx >= num_bad_words: |
| 125 | return |
| 126 | |
| 127 | pos = tl.load(expanded_local_pos_ptr + token_idx) |
| 128 | cur_req_first_pos = token_idx - pos |
| 129 | |
| 130 | prompt_len = tl.load(prompt_len_ptr + req_state_idx) |
| 131 | total_len = tl.load(total_len_ptr + req_state_idx) |
| 132 | output_len = total_len - prompt_len |
| 133 | effective_len = output_len + pos |
| 134 | |
| 135 | bd_offsets_base = bad_word_offsets_ptr + req_state_idx * bad_word_offsets_stride |
| 136 | bd_tokens_base = bad_word_token_ids_ptr + req_state_idx * bad_word_token_ids_stride |
| 137 | output_base = all_token_ids_ptr + req_state_idx * all_token_ids_stride + prompt_len |
| 138 | |
| 139 | start = tl.load(bd_offsets_base + bw_idx) |
| 140 | end = tl.load(bd_offsets_base + bw_idx + 1) |
| 141 | bad_word_len = end - start |
| 142 | prefix_len = bad_word_len - 1 |
| 143 | |
| 144 | if prefix_len > effective_len: |
| 145 | return |
| 146 | |
| 147 | last_token = tl.load(bd_tokens_base + end - 1) |
| 148 | match = 1 |
| 149 | for i in range(prefix_len): |
| 150 | expected = tl.load(bd_tokens_base + start + i) |
| 151 | actual_pos = effective_len - prefix_len + i |
| 152 | |
| 153 | from_spec_input = actual_pos >= output_len |
| 154 | if from_spec_input: |
| 155 | # input_ids at local position 0 is the last committed token; |
| 156 | # draft tokens start at local position 1. |
| 157 | spec_offset = actual_pos - output_len |
| 158 | actual = tl.load(input_ids_ptr + cur_req_first_pos + spec_offset + 1) |
| 159 | else: |
| 160 | actual = tl.load(output_base + actual_pos) |
| 161 | |
| 162 | match = match & (expected == actual) |
| 163 | |
| 164 | if match: |
| 165 | tl.store(logits_ptr + token_idx * logits_stride + last_token, -float("inf")) |
| 166 | |
| 167 | |
| 168 | def apply_bad_words( |
| 169 | logits: torch.Tensor, |
| 170 | expanded_idx_mapping: torch.Tensor, |
| 171 | bad_word_token_ids: torch.Tensor, |
| 172 | bad_word_offsets: torch.Tensor, |
| 173 | num_bad_words: torch.Tensor, |
| 174 | all_token_ids: torch.Tensor, |
| 175 | prompt_len: torch.Tensor, |
| 176 | total_len: torch.Tensor, |
| 177 | input_ids: torch.Tensor, |
| 178 | expanded_local_pos: torch.Tensor, |
| 179 | max_num_bad_words: int, |
| 180 | ) -> None: |
| 181 | num_tokens = logits.shape[0] |
| 182 | _bad_words_kernel[(num_tokens, max_num_bad_words)]( |
| 183 | logits, |
| 184 | logits.stride(0), |
| 185 | expanded_idx_mapping, |
| 186 | bad_word_token_ids, |
| 187 | bad_word_token_ids.stride(0), |
| 188 | bad_word_offsets, |
| 189 | bad_word_offsets.stride(0), |
| 190 | num_bad_words, |
| 191 | all_token_ids, |
| 192 | all_token_ids.stride(0), |
| 193 | prompt_len, |
| 194 | total_len, |
| 195 | input_ids, |
| 196 | expanded_local_pos, |
| 197 | ) |