| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import torch |
| 5 | |
| 6 | _SMALLEST_LOGIT = float("-inf") |
| 7 | |
| 8 | |
| 9 | def _apply_bad_words_single_batch( |
| 10 | logits: torch.Tensor, |
| 11 | bad_words_token_ids: list[list[int]], |
| 12 | past_tokens_ids: list[int], |
| 13 | ) -> None: |
| 14 | for bad_word_ids in bad_words_token_ids: |
| 15 | if len(bad_word_ids) > len(past_tokens_ids) + 1: |
| 16 | continue |
| 17 | |
| 18 | prefix_length = len(bad_word_ids) - 1 |
| 19 | last_token_id = bad_word_ids[-1] |
| 20 | actual_prefix = past_tokens_ids[-prefix_length:] if prefix_length > 0 else [] |
| 21 | expected_prefix = bad_word_ids[:prefix_length] |
| 22 | |
| 23 | assert len(actual_prefix) == len(expected_prefix) |
| 24 | |
| 25 | if actual_prefix == expected_prefix: |
| 26 | # Assign to slice to avoid cpu->gpu sync. |
| 27 | logits[last_token_id : last_token_id + 1] = _SMALLEST_LOGIT |
| 28 | |
| 29 | |
| 30 | def apply_bad_words( |
| 31 | logits: torch.Tensor, |
| 32 | bad_words_token_ids: dict[int, list[list[int]]], |
| 33 | past_tokens_ids: list[list[int]], |
| 34 | ) -> None: |
| 35 | for i, bad_words_ids in bad_words_token_ids.items(): |
| 36 | _apply_bad_words_single_batch(logits[i], bad_words_ids, past_tokens_ids[i]) |
| 37 | |
| 38 | |
| 39 | def apply_bad_words_with_drafts( |
| 40 | logits: torch.Tensor, |
| 41 | bad_words_token_ids: dict[int, list[list[int]]], |
| 42 | past_tokens_ids: list[list[int]], |
| 43 | num_draft_tokens: list[int], |
| 44 | ) -> None: |
| 45 | start_idx = 0 |
| 46 | remaining = len(bad_words_token_ids) |
| 47 | for i, n in enumerate(num_draft_tokens): |
| 48 | if (bad_words_ids := bad_words_token_ids.get(i)) is not None: |
| 49 | for draft_idx in range(start_idx, start_idx + n): |
| 50 | _apply_bad_words_single_batch( |
| 51 | logits[draft_idx], |
| 52 | bad_words_ids, |
| 53 | past_tokens_ids[draft_idx], |
| 54 | ) |
| 55 | remaining -= 1 |
| 56 | if not remaining: |
| 57 | break |
| 58 | start_idx += n |