| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterator |
| 5 | from enum import Enum |
| 6 | from typing import NamedTuple |
| 7 | from unittest.mock import MagicMock |
| 8 | |
| 9 | import regex as re |
| 10 | import torch |
| 11 | |
| 12 | from vllm import CompletionOutput |
| 13 | from vllm.config import ReasoningConfig |
| 14 | from vllm.utils.torch_utils import make_tensor_with_pad |
| 15 | from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor |
| 16 | from vllm.v1.sample.metadata import SamplingMetadata |
| 17 | |
| 18 | |
| 19 | def create_mock_reasoning_config( |
| 20 | start_token_ids: list[int], end_token_ids: list[int] |
| 21 | ) -> ReasoningConfig: |
| 22 | return MagicMock( |
| 23 | spec=ReasoningConfig, |
| 24 | reasoning_start_token_ids=start_token_ids, |
| 25 | reasoning_end_token_ids=end_token_ids, |
| 26 | ) |
| 27 | |
| 28 | |
| 29 | class BatchLogprobsComposition(Enum): |
| 30 | """Types of logprobs configs to include in test batch.""" |
| 31 | |
| 32 | NONE = 0 |
| 33 | SAMPLE = 1 |
| 34 | PROMPT = 2 |
| 35 | SAMPLE_PROMPT = 3 |
| 36 | |
| 37 | |
| 38 | BatchLogprobsSpecType = list[tuple[int | None, int | None]] |
| 39 | |
| 40 | |
| 41 | def get_test_batch( |
| 42 | batch_logprobs_composition: BatchLogprobsComposition, |
| 43 | ) -> BatchLogprobsSpecType: |
| 44 | """Generate logprobs configs for a batch of requests. |
| 45 | |
| 46 | A given request's logprobs configuration is (1) num_sample_logprobs and (2) |
| 47 | num_prompt_logprobs. The batch logprobs configuration is the list of request |
| 48 | logprobs configs. |
| 49 | |
| 50 | batch_logprobs_composition == NONE yields a batch with no sample or prompt |
| 51 | logprobs |
| 52 | |
| 53 | batch_logprobs_composition == SAMPLE yields a batch with some requests |
| 54 | configured for sample logprobs only, and others configured for no logprobs |
| 55 | |
| 56 | batch_logprobs_composition == PROMPT yields a batch with some requests |
| 57 | configured for prompt logprobs only, and others configured for no logprobs |
| 58 | |
| 59 | batch_logprobs_composition == SAMPLE_PROMPT yields a batch with some |
| 60 | requests configured for sample logprobs and prompt logprobs, some configured |
| 61 | for only sample logprobs or only prompt logprobs, and some configured for |
| 62 | no logprobs |
| 63 | |
| 64 | Args: |
| 65 | batch_logprobs_composition: types of logprobs configs to include in batch |
| 66 | |
| 67 | Returns: |
| 68 | list of (Optional[num_sample_logprobs], Optional[num_prompt_logprobs]) |
| 69 | tuples |
| 70 | |
| 71 | """ |
| 72 | if batch_logprobs_composition == BatchLogprobsComposition.NONE: |
| 73 | # No requests with sample or prompt logprobs |
| 74 | return [(None, None)] |
| 75 | elif batch_logprobs_composition == BatchLogprobsComposition.SAMPLE: |
| 76 | # Requests requiring sample logprobs or no logprobs |
| 77 | return [ |
| 78 | (None, None), |
| 79 | (0, None), |
| 80 | (5, None), |
| 81 | (3, None), |
| 82 | ] |
| 83 | elif batch_logprobs_composition == BatchLogprobsComposition.PROMPT: |
| 84 | # Requests requiring prompt logprobs or no logprobs |
| 85 | return [ |
| 86 | (None, None), |
| 87 | (None, 0), |
| 88 | (None, 6), |
| 89 | (None, 5), |
| 90 | ] |
| 91 | elif batch_logprobs_composition == BatchLogprobsComposition.SAMPLE_PROMPT: |
| 92 | # Requests requiring either no logprobs, just |
| 93 | # sample logprobs, just prompt logprobs, or |
| 94 | # both sample and prompt logprobs |
| 95 | return [ |
| 96 | (None, None), |
| 97 | (0, None), |
| 98 | (5, None), |
| 99 | (3, None), |
| 100 | (0, 3), |
| 101 | (6, 0), |
| 102 | (6, 3), |
| 103 | (None, 6), |
| 104 | (None, 5), |
| 105 | (None, 0), |
| 106 | ] |
| 107 | else: |
| 108 | raise ValueError("Invalid logprobs batch configuration for test.") |
| 109 | |
| 110 | |
| 111 | def assert_incr_detok_str_matches_non_incr_detok_str( |
| 112 | incremental_detokenization_str: str, |
| 113 | non_incremental_detokenization_str: str, |
| 114 | msg: str, |
| 115 | ) -> None: |
| 116 | """Compare incrementally detok. text to non-incrementally detok. text. |
| 117 | |
| 118 | Fail if the strings mismatch after non-alphanumeric characters are stripped |
| 119 | out. |
| 120 | |
| 121 | Rationale: incremental detokenization in the text generation process allows |
| 122 | the tokenizer to adjust the next token text output based on the token's |
| 123 | context in the string. However, logprobs detokenization detokenizes each |
| 124 | token individually, and the resultant strings may include some |
| 125 | non-alphanumeric placeholder characters where there could be i.e. |
| 126 | whitespace. So, this function compares only the alphanumeric text |
| 127 | between two strings and fails if there is a mismatch, which helps |
| 128 | with validating logprobs detokenization. |
| 129 | |
| 130 | Args: |
| 131 | incremental_detokenization_str: incrementally-detokenized generated text |
| 132 | non_incremental_detokenization_str: non-incrementally-detokenized logprob |
| 133 | tokens |
| 134 | msg: error message if `assert` fails |
| 135 | |
| 136 | """ |
| 137 | rgx = r"[^a-zA-Z0-9]+" |
| 138 | assert re.sub(rgx, "", incremental_detokenization_str) == re.sub( |
| 139 | rgx, "", non_incremental_detokenization_str |
| 140 | ), msg |
| 141 | |
| 142 | |
| 143 | def compute_correct_cumulative_logprob(completion_output: CompletionOutput) -> float: |
| 144 | """Compute known-good value for evaluating cumulative logprob. |
| 145 | |
| 146 | Args: |
| 147 | completion_output: completion output from engine |
| 148 | |
| 149 | Returns: |
| 150 | Known-good cumulative logprob value |
| 151 | |
| 152 | """ |
| 153 | token_ids = completion_output.token_ids |
| 154 | logprobs = completion_output.logprobs |
| 155 | assert logprobs is not None |
| 156 | values: list[float] = [] |
| 157 | for tok_id, lp in zip(token_ids, logprobs): |
| 158 | assert lp is not None |
| 159 | values.append(lp[tok_id].logprob) |
| 160 | return sum(values) |
| 161 | |
| 162 | |
| 163 | def create_fake_logits(batch_size: int, vocab_size: int) -> torch.Tensor: |
| 164 | fake_logits = torch.full((batch_size, vocab_size), 1e-2, dtype=torch.float) |
| 165 | return fake_logits |
| 166 | |
| 167 | |
| 168 | def create_penalty_tensor( |
| 169 | batch_size: int, penalty_value: float, device: torch.device |
| 170 | ) -> torch.Tensor: |
| 171 | return torch.full( |
| 172 | (batch_size,), fill_value=penalty_value, dtype=torch.float, device=device |
| 173 | ) |
| 174 | |
| 175 | |
| 176 | def create_prompt_tokens_tensor( |
| 177 | prompt_token_ids: list[list[int]], |
| 178 | vocab_size: int, |
| 179 | device: torch.device, |
| 180 | ) -> torch.Tensor: |
| 181 | return make_tensor_with_pad( |
| 182 | prompt_token_ids, |
| 183 | pad=vocab_size, |
| 184 | device=device, |
| 185 | dtype=torch.int64, |
| 186 | pin_memory=False, |
| 187 | ) |
| 188 | |
| 189 | |
| 190 | class LogitsprocsTestFakes(NamedTuple): |
| 191 | """Wraps fake data structures to support testing.""" |
| 192 | |
| 193 | logits: torch.Tensor |
| 194 | sampling_metadata: SamplingMetadata |
| 195 | |
| 196 | def get_logitsprocs_by_cls( |
| 197 | self, |
| 198 | cls: type[LogitsProcessor], |
| 199 | ) -> Iterator[LogitsProcessor]: |
| 200 | """Yield logits processors of a specific class. |
| 201 | |
| 202 | Args: |
| 203 | cls: :class:`LogitsProcessor` subclass |
| 204 | |
| 205 | Returns: |
| 206 | Iterator over logits processors |
| 207 | |
| 208 | """ |
| 209 | return ( |
| 210 | lp for lp in self.sampling_metadata.logitsprocs.all if isinstance(lp, cls) |
| 211 | ) |
| 212 | |
| 213 | def get_logitsprocs(self) -> Iterator[LogitsProcessor]: |
| 214 | """Iterator over all logits processors.""" |
| 215 | return self.sampling_metadata.logitsprocs.all |
| 216 | |
| 217 | |
| 218 | def fake_update_logitsprocs_state( |
| 219 | test_fakes: LogitsprocsTestFakes, |
| 220 | batch_update: BatchUpdate | None, |
| 221 | ) -> None: |
| 222 | """Imitate logits processors persistent batch state update |
| 223 | in engine core""" |
| 224 | for logitproc in test_fakes.get_logitsprocs(): |
| 225 | logitproc.update_state(batch_update) |
| 226 | holder = test_fakes.sampling_metadata.thinking_budget_state_holder |
| 227 | if holder is not None: |
| 228 | holder.sync_batch(batch_update) |
| 229 | |
| 230 | |
| 231 | def fake_apply_logitsprocs( |
| 232 | test_fakes: LogitsprocsTestFakes, |
| 233 | slice_indices: list[int], |
| 234 | slot_output_token_ids: list[list[int]] | None = None, |
| 235 | ) -> torch.Tensor: |
| 236 | """Imitate application of logits processors in engine core. |
| 237 | |
| 238 | When ``thinking_budget_state_holder`` has tracked requests, this mirrors |
| 239 | :meth:`Sampler.apply_logits_processors` by refreshing per-slot |
| 240 | ``output_token_ids`` (if ``slot_output_token_ids`` is provided), then |
| 241 | ``update_state`` + ``apply_to_logits`` on the holder after built-in logits |
| 242 | processors. |
| 243 | """ |
| 244 | logits = test_fakes.logits[torch.tensor(slice_indices, dtype=torch.long)].clone() |
| 245 | for processor in test_fakes.get_logitsprocs(): |
| 246 | logits = processor.apply(logits) |
| 247 | |
| 248 | md = test_fakes.sampling_metadata |
| 249 | holder = md.thinking_budget_state_holder |
| 250 | if holder is not None and holder.has_tracked_requests(): |
| 251 | if slot_output_token_ids is not None: |
| 252 | for i, toks in enumerate(slot_output_token_ids): |
| 253 | if i < len(md.output_token_ids): |
| 254 | md.output_token_ids[i] = list(toks) |
| 255 | holder.update_state(md.output_token_ids, md.spec_token_ids, None) |
| 256 | logits = holder.apply_to_logits(logits, False, md.spec_token_ids) |
| 257 | return logits |
| 258 | |
| 259 | |
| 260 | def create_allowed_token_ids( |
| 261 | batch_size: int, |
| 262 | vocab_size: int, |
| 263 | num_allowed_token_ids: int, |
| 264 | device: torch.device, |
| 265 | ) -> torch.Tensor | None: |
| 266 | mask: torch.Tensor | None = None |
| 267 | for i in range(batch_size): |
| 268 | if i % 2 == 1: |
| 269 | continue |
| 270 | if mask is None: |
| 271 | mask = torch.zeros( |
| 272 | (batch_size, vocab_size), dtype=torch.bool, device=device |
| 273 | ) |
| 274 | start = min(i, vocab_size - 1) |
| 275 | end = min(i + num_allowed_token_ids, vocab_size - 1) |
| 276 | mask[i, start:end] = True |
| 277 | return mask |