@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / eb8b5eb183

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 19:45 UTC

[V1] Support bad_words in sampler (#13376)

Signed-off-by: 22quinn <33176974+22quinn@users.noreply.github.com> Co-authored-by: Nick Hill <nhill@redhat.com>

22quinn committed Mar 8, 2025 at 14:50 UTC eb8b5eb183b8428f7e58adf2559e7a8d9400fc30

13 files changed +266 -28

tests/test_utils.py

+24 -1

@@ -14,7 +14,7 @@ from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config

14 from vllm.utils import (FlexibleArgumentParser, MemorySnapshot,

15 PlaceholderModule, StoreBoolean, bind_kv_cache,

16 deprecate_kwargs, get_open_port, memory_profiling,

17 - merge_async_iterators, supports_kw)

17 + merge_async_iterators, supports_kw, swap_dict_values)

19 from .utils import error_on_warning, fork_new_process_for_each_test

@@ -449,3 +449,26 @@ def test_placeholder_module_error_handling():

449 with build_ctx():

450 # Test conflict with internal __module attribute

451 _ = placeholder_attr.module

452 +

453 +

454 +@pytest.mark.parametrize(

455 + "obj,key1,key2",

456 + [

457 + # Tests for both keys exist

458 + ({1: "a", 2: "b"}, 1, 2),

459 + # Tests for one key does not exist

460 + ({1: "a", 2: "b"}, 1, 3),

461 + # Tests for both keys do not exist

462 + ({1: "a", 2: "b"}, 3, 4),

463 + ])

464 +def test_swap_dict_values(obj, key1, key2):

465 + original_obj = obj.copy()

466 + swap_dict_values(obj, key1, key2)

467 + if key1 in original_obj:

468 + assert obj[key2] == original_obj[key1]

469 + else:

470 + assert key2 not in obj

471 + if key2 in original_obj:

472 + assert obj[key1] == original_obj[key2]

473 + else:

474 + assert key1 not in obj

tests/v1/sample/test_rejection_sampler.py

@@ -42,6 +42,7 @@ def create_sampling_metadata(spec_tokens: list[list[int]]) -> SamplingMetadata:

42 min_tokens={},

43 logit_bias=[None] * batch_size,

44 allowed_token_ids_mask=None,

45 + bad_words_token_ids={},

46 )

tests/v1/sample/test_sampler.py

+76

@@ -77,6 +77,49 @@ def _create_allowed_token_ids(

77 return mask

80 +def _create_bad_words_token_ids(

81 + batch_size: int, vocab_size: int,

82 + bad_words_lengths: list[tuple[int]]) -> dict[int, list[list[int]]]:

83 + bad_words_token_ids = {}

84 + for batch_idx in range(batch_size):

85 + token_ids_single_batch = []

86 + for bad_words_length in bad_words_lengths:

87 + token_ids = np.random.choice(vocab_size,

88 + size=bad_words_length,

89 + replace=True).tolist()

90 + token_ids_single_batch.append(token_ids)

91 + bad_words_token_ids[batch_idx] = token_ids_single_batch

92 + if batch_size >= 2:

93 + # Test no bad_words for some batch

94 + no_bad_words_batch_idx = np.random.choice(batch_size)

95 + bad_words_token_ids.pop(no_bad_words_batch_idx, None)

96 + return bad_words_token_ids

97 +

98 +

99 +def _update_output_token_ids_for_bad_words(

100 + metadata: SamplingMetadata, vocab_size: int) -> dict[int, list[int]]:

101 + bad_words_last_tokens = {}

102 + for batch_idx, bad_words_token_ids in metadata.bad_words_token_ids.items():

103 + output_token_ids = metadata.output_token_ids[batch_idx]

104 + bad_words_last_token: list[int] = []

105 + for i, bad_word_token_ids in enumerate(bad_words_token_ids):

106 + if len(bad_word_token_ids) == 1:

107 + # Single token id always affects logits

108 + bad_words_last_token.append(bad_word_token_ids[0])

109 + else:

110 + prefix_length = len(bad_word_token_ids) - 1

111 + has_bad_words = np.random.choice([True, False])

112 + if has_bad_words:

113 + output_token_ids[-prefix_length:] = bad_word_token_ids[:-1]

114 + bad_words_last_token.append(bad_word_token_ids[-1])

115 + break # Maximum one update to output_token_ids

116 + else: # Make sure no accidental match to bad words

117 + output_token_ids[-1] = (bad_word_token_ids[-2] +

118 + 1) % vocab_size

119 + bad_words_last_tokens[batch_idx] = bad_words_last_token

120 + return bad_words_last_tokens

121 +

122 +

123 def _create_default_sampling_metadata(

124 num_output_tokens: int,

125 batch_size: int,

@@ -112,6 +155,7 @@ def _create_default_sampling_metadata(

155 min_tokens={},

156 logit_bias=[None] * batch_size,

157 allowed_token_ids_mask=None,

158 + bad_words_token_ids={},

159 )

160 return fake_sampling_metadata

161

@@ -467,3 +511,35 @@ def test_sampler_allowed_token_ids(device: str, batch_size: int,

511 "inf"), f"{batch_idx}, {token_id}"

512 else:

513 assert logits_for_req[token_id] != -float("inf")

514 +

515 +

516 +@pytest.mark.parametrize("device", CUDA_DEVICES)

517 +@pytest.mark.parametrize("batch_size", [1, 2, 32])

518 +@pytest.mark.parametrize("bad_words_lengths", [(1, ), (1, 3), (2, 2)])

519 +def test_sampler_bad_words(device: str, batch_size: int,

520 + bad_words_lengths: list[tuple[int]]):

521 + """

522 + Test to verify that when the bad words restriction is present, tokens

523 + are penalized based on their match with the bad words.

524 + """

525 + torch.set_default_device(device)

526 + # Create fake logits where each token is assigned the same

527 + # logit value.

528 + fake_logits = _create_fake_logits(batch_size, VOCAB_SIZE)

529 + sampling_metadata = _create_default_sampling_metadata(

530 + NUM_OUTPUT_TOKENS, batch_size, VOCAB_SIZE, torch.device(device))

531 + sampling_metadata.bad_words_token_ids = _create_bad_words_token_ids(

532 + batch_size, VOCAB_SIZE, bad_words_lengths)

533 + bad_words_last_tokens = _update_output_token_ids_for_bad_words(

534 + sampling_metadata, VOCAB_SIZE)

535 + sampler = Sampler()

536 + logits = sampler.apply_bad_words(fake_logits, sampling_metadata)

537 + logits = logits.cpu()

538 + for batch_idx in range(batch_size):

539 + logits_for_req = logits[batch_idx]

540 + for token_id in range(VOCAB_SIZE):

541 + if (batch_idx in bad_words_last_tokens

542 + and token_id in bad_words_last_tokens[batch_idx]):

543 + assert logits_for_req[token_id] == -float("inf")

544 + else:

545 + assert logits_for_req[token_id] != -float("inf")

tests/v1/sample/test_sampling_params_e2e.py

+16 -2

@@ -120,8 +120,22 @@ def test_detokenize_false(model):

120 def test_bad_words(model):

121 """Check that we respect bad words."""

122

123 - with pytest.raises(ValueError):

124 - _ = model.generate(PROMPT, SamplingParams(bad_words=["Hello"]))

123 + output = model.generate(PROMPT, SamplingParams(temperature=0))

124 + split_text = output[0].outputs[0].text.split()

125 +

126 + bad_words_1 = " ".join(split_text[:2])

127 + params = SamplingParams(temperature=0, bad_words=[bad_words_1])

128 + output = model.generate(PROMPT, params)

129 + new_text = output[0].outputs[0].text

130 + assert bad_words_1 not in new_text

131 +

132 + bad_words_2 = new_text.split()[-1]

133 + params = SamplingParams(temperature=0,

134 + bad_words=[bad_words_1, bad_words_2])

135 + output = model.generate(PROMPT, params)

136 + new_text = output[0].outputs[0].text

137 + assert bad_words_1 not in new_text

138 + assert bad_words_2 not in new_text

139

140

141 def test_logits_processor(model):

tests/v1/worker/test_gpu_input_batch.py

@@ -100,6 +100,7 @@ def _construct_expected_sampling_metadata(

100 VOCAB_SIZE,

101 dtype=torch.bool,

102 device=device)

103 + bad_words_token_ids = {}

104 for req in reqs:

105 if req.req_id not in req_ids_retained:

106 continue

@@ -123,6 +124,8 @@ def _construct_expected_sampling_metadata(

124 if req.sampling_params.allowed_token_ids:

125 allowed_token_ids_mask[index_in_input_batch][

126 req.sampling_params.allowed_token_ids] = True

127 + bad_words_token_ids[

128 + index_in_input_batch] = req.sampling_params.bad_words_token_ids

129

130 return SamplingMetadata(

131 temperature=torch.tensor(temperature, dtype=torch.float,

@@ -159,6 +162,7 @@ def _construct_expected_sampling_metadata(

162 and all(x == 1 for x in repetition_penalties)),

163 logit_bias=logit_bias,

164 allowed_token_ids_mask=allowed_token_ids_mask,

165 + bad_words_token_ids=bad_words_token_ids,

166 )

167

168

@@ -284,6 +288,8 @@ def test_sampling_metadata_in_input_batch(device: str, batch_size: int):

288 assert torch.allclose(

289 expected_sampling_metadata.allowed_token_ids_mask,

290 sampling_metadata.allowed_token_ids_mask)

291 + assert expected_sampling_metadata.bad_words_token_ids == \

292 + sampling_metadata.bad_words_token_ids

293

294

295 @pytest.mark.parametrize("device", CUDA_DEVICES)

vllm/sampling_params.py

+51 -1

@@ -11,6 +11,8 @@ from pydantic import BaseModel

12 from vllm.logger import init_logger

13 from vllm.logits_process import LogitsProcessor

14 +from vllm.transformers_utils.tokenizer import AnyTokenizer

15 +from vllm.transformers_utils.tokenizers.mistral import MistralTokenizer

17 logger = init_logger(__name__)

@@ -202,7 +204,6 @@ class SamplingParams(

204 seed: Optional[int] = None

205 stop: Optional[Union[str, list[str]]] = None

206 stop_token_ids: Optional[list[int]] = None

205 - bad_words: Optional[list[str]] = None

207 ignore_eos: bool = False

208 max_tokens: Optional[int] = 16

209 min_tokens: int = 0

@@ -232,6 +233,10 @@ class SamplingParams(

233 allowed_token_ids: Optional[list[int]] = None

234 extra_args: Optional[dict[str, Any]] = None

235

236 + # Fields used for bad words

237 + bad_words: Optional[list[str]] = None

238 + _bad_words_token_ids: list[list[int]] = msgspec.field(default_factory=list)

239 +

240 @staticmethod

241 def from_optional(

242 n: Optional[int] = 1,

@@ -464,6 +469,46 @@ class SamplingParams(

469 eos_ids.update(self.stop_token_ids)

470 self.stop_token_ids = list(eos_ids)

471

472 + def update_from_tokenizer(self, tokenizer: AnyTokenizer) -> None:

473 + if self.bad_words is None:

474 + return

475 + for bad_word in self.bad_words:

476 + # To prohibit words both at the beginning

477 + # and in the middle of text

478 + # (related to add_prefix_space tokenizer parameter)

479 + for add_prefix_space in [False, True]:

480 + prefix = " " if add_prefix_space else ""

481 + prompt = prefix + bad_word.lstrip()

482 +

483 + if isinstance(tokenizer, MistralTokenizer):

484 + # Mistral tokenizers should not add special tokens

485 + prompt_token_ids = tokenizer.encode(text=prompt)

486 + else:

487 + prompt_token_ids = tokenizer.encode(

488 + text=prompt, add_special_tokens=False)

489 +

490 + # If no space at the beginning

491 + # or if prefix space produces a new word token

492 + if (not add_prefix_space) or (

493 + add_prefix_space and prompt_token_ids[0]

494 + != self._bad_words_token_ids[-1][0]

495 + and len(prompt_token_ids) == len(

496 + self._bad_words_token_ids[-1])):

497 + self._bad_words_token_ids.append(prompt_token_ids)

498 +

499 + invalid_token_ids = [

500 + token_id for bad_words_token_ids in self._bad_words_token_ids

501 + for token_id in bad_words_token_ids

502 + if token_id < 0 or token_id > tokenizer.max_token_id

503 + ]

504 + if len(invalid_token_ids) > 0:

505 + raise ValueError(

506 + f"The model vocabulary size is {tokenizer.max_token_id+1},"

507 + f" but the following tokens"

508 + f" were specified as bad: {invalid_token_ids}."

509 + f" All token id values should be integers satisfying:"

510 + f" 0 <= token_id <= {tokenizer.max_token_id}.")

511 +

512 @cached_property

513 def sampling_type(self) -> SamplingType:

514 if self.temperature < _SAMPLING_EPS:

@@ -476,6 +521,11 @@ class SamplingParams(

521 def all_stop_token_ids(self) -> set[int]:

522 return self._all_stop_token_ids

523

524 + @property

525 + def bad_words_token_ids(self) -> list[list[int]]:

526 + # For internal use only. Backward compatibility not guaranteed

527 + return self._bad_words_token_ids

528 +

529 def clone(self) -> "SamplingParams":

530 """Deep copy, but maybe not the LogitsProcessor objects.

531

vllm/utils.py

+16

@@ -2361,3 +2361,19 @@ class LazyLoader(types.ModuleType):

2361 if self._module is None:

2362 self._module = self._load()

2363 return dir(self._module)

2364 +

2365 +

2366 +def swap_dict_values(obj: dict[_K, _V], key1: _K, key2: _K) -> None:

2367 + """

2368 + Helper function to swap values for two keys

2369 + """

2370 + v1 = obj.get(key1)

2371 + v2 = obj.get(key2)

2372 + if v1 is not None:

2373 + obj[key2] = v1

2374 + else:

2375 + obj.pop(key2, None)

2376 + if v2 is not None:

2377 + obj[key1] = v2

2378 + else:

2379 + obj.pop(key1, None)

vllm/v1/engine/processor.py

+2 -3

@@ -94,9 +94,6 @@ class Processor:

94 # Best of not yet supported.

95 if params.best_of is not None and params.best_of > 1:

96 raise ValueError("VLLM V1 does not yet support best_of.")

97 - # Bad words not yet supported.

98 - if params.bad_words:

99 - raise ValueError("VLLM V1 does not yet support bad_words.")

97 # Logits processors not supported.

98 if params.logits_processors:

99 raise ValueError("VLLM V1 does not support per request "

@@ -203,6 +200,8 @@ class Processor:

200 sampling_params = params.clone()

201 sampling_params.update_from_generation_config(

202 self.generation_config_fields, eos_token_id)

203 + sampling_params.update_from_tokenizer(

204 + self.tokenizer.get_lora_tokenizer(lora_request))

205

206 # Multimodal related.

207 # Compute MM hashes (if enabled)

vllm/v1/sample/metadata.py

@@ -38,3 +38,6 @@ class SamplingMetadata:

38 # `allowed_token_ids_mask` is a 2D bool tensor of shape (max batch size,

39 # vocab size).

40 allowed_token_ids_mask: Optional[torch.Tensor]

41 +

42 + # req_index -> bad_words_token_ids

43 + bad_words_token_ids: dict[int, list[list[int]]]

vllm/v1/sample/ops/bad_words.py new

+38

@@ -0,0 +1,38 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +

3 +import torch

4 +

5 +_SMALLEST_LOGIT = float("-inf")

6 +

7 +

8 +def _apply_bad_words_single_batch(

9 + logits: torch.Tensor,

10 + bad_words_token_ids: list[list[int]],

11 + past_tokens_ids: list[int],

12 +) -> None:

13 + for bad_word_ids in bad_words_token_ids:

14 + if len(bad_word_ids) > len(past_tokens_ids) + 1:

15 + continue

16 +

17 + prefix_length = len(bad_word_ids) - 1

18 + last_token_id = bad_word_ids[-1]

19 + if prefix_length > 0:

20 + actual_prefix = past_tokens_ids[-prefix_length:]

21 + else:

22 + actual_prefix = []

23 + expected_prefix = bad_word_ids[:prefix_length]

24 +

25 + assert len(actual_prefix) == len(expected_prefix)

26 +

27 + if actual_prefix == expected_prefix:

28 + logits[last_token_id] = _SMALLEST_LOGIT

29 +

30 +

31 +def apply_bad_words(

32 + logits: torch.Tensor,

33 + bad_words_token_ids: dict[int, list[list[int]]],

34 + past_tokens_ids: list[list[int]],

35 +) -> None:

36 + for i, bad_words_ids in bad_words_token_ids.items():

37 + _apply_bad_words_single_batch(logits[i], bad_words_ids,

38 + past_tokens_ids[i])

vllm/v1/sample/sampler.py

+16

@@ -6,6 +6,7 @@ import torch.nn as nn

7 from vllm.v1.outputs import LogprobsTensors, SamplerOutput

8 from vllm.v1.sample.metadata import SamplingMetadata

9 +from vllm.v1.sample.ops.bad_words import apply_bad_words

10 from vllm.v1.sample.ops.penalties import (apply_all_penalties,

11 apply_min_token_penalties)

12 from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler

@@ -38,6 +39,8 @@ class Sampler(nn.Module):

39 logits = logits.to(torch.float32)

40 # Apply allowed token ids.

41 logits = self.apply_allowed_token_ids(logits, sampling_metadata)

42 + # Apply bad words exclusion.

43 + logits = self.apply_bad_words(logits, sampling_metadata)

44 # Apply logits bias.

45 logits = self.apply_logits_bias(logits, sampling_metadata)

46 # Apply penalties (e.g., min_tokens, freq_penalties).

@@ -237,3 +240,16 @@ class Sampler(nn.Module):

240 logits.masked_fill_(sampling_metadata.allowed_token_ids_mask,

241 float("-inf"))

242 return logits

243 +

244 + def apply_bad_words(

245 + self,

246 + logits: torch.Tensor,

247 + sampling_metadata: SamplingMetadata,

248 + ) -> torch.Tensor:

249 + if sampling_metadata.bad_words_token_ids:

250 + apply_bad_words(

251 + logits,

252 + sampling_metadata.bad_words_token_ids,

253 + sampling_metadata.output_token_ids,

254 + )

255 + return logits

vllm/v1/worker/gpu_input_batch.py

+16 -21

@@ -10,6 +10,7 @@ import torch

10 from vllm.lora.request import LoRARequest

11 from vllm.multimodal import MultiModalKwargs

12 from vllm.sampling_params import SamplingParams, SamplingType

13 +from vllm.utils import swap_dict_values

14 from vllm.v1.sample.metadata import SamplingMetadata

15 from vllm.v1.utils import copy_slice

16 from vllm.v1.worker.block_table import BlockTable

@@ -204,6 +205,9 @@ class InputBatch:

205 self.allowed_token_ids_mask: Optional[torch.Tensor] = None

206 self.allowed_token_ids_mask_cpu_tensor: Optional[torch.Tensor] = None

207

208 + # req_index -> bad_words_token_ids

209 + self.bad_words_token_ids: dict[int, list[list[int]]] = {}

210 +

211 self.req_output_token_ids: list[Optional[list[int]]] = []

212

213 # This is updated each time the batch constituents change.

@@ -320,6 +324,9 @@ class InputBatch:

324 self.allowed_token_ids_mask_cpu_tensor[req_index][

325 sampling_params.allowed_token_ids] = False

326

327 + self.bad_words_token_ids[

328 + req_index] = sampling_params.bad_words_token_ids

329 +

330 # Add request lora ID

331 if request.lora_request:

332 lora_id = request.lora_request.lora_int_id

@@ -369,6 +376,7 @@ class InputBatch:

376 if self.allowed_token_ids_mask_cpu_tensor is not None:

377 # False means we don't fill with -inf.

378 self.allowed_token_ids_mask_cpu_tensor[req_index].fill_(False)

379 + self.bad_words_token_ids.pop(req_index, None)

380 return req_index

381

382 def swap_states(self, i1: int, i2: int) -> None:

@@ -413,27 +421,9 @@ class InputBatch:

421 self.token_ids_cpu[i1, ...] = self.token_ids_cpu[i2, ...]

422 self.token_ids_cpu[i2, ...] = tmp

423

416 - g1 = self.generators.get(i1)

417 - g2 = self.generators.get(i2)

418 - if g1 is not None:

419 - self.generators[i2] = g1

420 - else:

421 - self.generators.pop(i2, None)

422 - if g2 is not None:

423 - self.generators[i1] = g2

424 - else:

425 - self.generators.pop(i1, None)

426 -

427 - t1 = self.min_tokens.get(i1)

428 - t2 = self.min_tokens.get(i2)

429 - if t1 is not None:

430 - self.min_tokens[i2] = t1

431 - else:

432 - self.min_tokens.pop(i2, None)

433 - if t2 is not None:

434 - self.min_tokens[i1] = t2

435 - else:

436 - self.min_tokens.pop(i1, None)

424 + swap_dict_values(self.generators, i1, i2)

425 + swap_dict_values(self.min_tokens, i1, i2)

426 + swap_dict_values(self.bad_words_token_ids, i1, i2)

427

428 self.request_lora_mapping[i1], self.request_lora_mapping[i2] =\

429 self.request_lora_mapping[i2], self.request_lora_mapping[i1]

@@ -518,6 +508,10 @@ class InputBatch:

508 empty_index] = self.allowed_token_ids_mask_cpu_tensor[

509 last_req_index]

510

511 + bad_words_token_ids = self.bad_words_token_ids.pop(

512 + last_req_index, None)

513 + if bad_words_token_ids is not None:

514 + self.bad_words_token_ids[empty_index] = bad_words_token_ids

515 # Decrement last_req_index since it is now empty.

516 last_req_index -= 1

517

@@ -585,6 +579,7 @@ class InputBatch:

579 no_penalties=self.no_penalties,

580 logit_bias=self.logit_bias[:num_reqs],

581 allowed_token_ids_mask=allowed_token_ids_mask,

582 + bad_words_token_ids=self.bad_words_token_ids,

583 )

584

585 def _make_prompt_token_ids_tensor(self) -> torch.Tensor:

vllm/v1/worker/gpu_model_runner.py

@@ -1268,6 +1268,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):

1268 min_tokens={},

1269 logit_bias=[None for _ in range(num_reqs)],

1270 allowed_token_ids_mask=None,

1271 + bad_words_token_ids={},

1272 )

1273 sampler_output = self.model.sample(logits=logits,

1274 sampling_metadata=dummy_metadata)