| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import random |
| 5 | from dataclasses import dataclass |
| 6 | from typing import TypeAlias |
| 7 | |
| 8 | import numpy as np |
| 9 | import torch |
| 10 | from transformers import PythonBackend, TokenizersBackend |
| 11 | |
| 12 | from vllm.engine.arg_utils import EngineArgs |
| 13 | from vllm.v1.engine import EngineCoreOutput, FinishReason |
| 14 | from vllm.v1.metrics.stats import PrefillStats |
| 15 | from vllm.v1.outputs import LogprobsLists, LogprobsTensors |
| 16 | |
| 17 | GeneralTokenizerType: TypeAlias = PythonBackend | TokenizersBackend |
| 18 | |
| 19 | # Number of sample logprobs to request when testing sample logprobs |
| 20 | NUM_SAMPLE_LOGPROBS_UNDER_TEST = 5 |
| 21 | # Number of prompt logprobs to request when testing prompt logprobs |
| 22 | NUM_PROMPT_LOGPROBS_UNDER_TEST = 7 |
| 23 | |
| 24 | TOKENIZER_NAME = "meta-llama/Llama-3.2-1B" |
| 25 | |
| 26 | FULL_STRINGS = [ |
| 27 | "My name is Robert from Neural Magic and I love working on vLLM so much!", |
| 28 | "Red Hat is the best open source company by far across Linux, K8s, and AI.", |
| 29 | "Nick is the name of my brother in addition to my colleague from Red Hat.", |
| 30 | ] |
| 31 | STOP_STRINGS = ["I love working on", "company by far", "brother in"] |
| 32 | PROMPT_LEN = 5 |
| 33 | |
| 34 | random.seed(42) |
| 35 | |
| 36 | |
| 37 | def _create_random_top_logprob_test_vector( |
| 38 | num_logprobs: int, |
| 39 | lower: float, |
| 40 | upper: float, |
| 41 | ) -> torch.Tensor: |
| 42 | """Create a random vector of top logprob float values. |
| 43 | |
| 44 | Use to create fake sample logprobs for testing. |
| 45 | |
| 46 | Note that a real production scenario would require |
| 47 | logprobs to be sorted in descending order, something |
| 48 | which is omitted in this function. |
| 49 | |
| 50 | Args: |
| 51 | num_logprobs: number of top logprobs |
| 52 | lower: lower range of logprob float values |
| 53 | upper: upper range of logprob float values |
| 54 | |
| 55 | Returns: |
| 56 | 1D length-`num_logprobs` torch Tensor of float logprob values |
| 57 | |
| 58 | """ |
| 59 | return torch.rand(num_logprobs) * (upper - lower) + lower |
| 60 | |
| 61 | |
| 62 | def _create_random_top_logprob_test_matrix( |
| 63 | shape: tuple, |
| 64 | lower: float, |
| 65 | upper: float, |
| 66 | ) -> torch.Tensor: |
| 67 | """Create a random matrix of top logprob float values. |
| 68 | |
| 69 | Use to create fake prompt logprobs for testing. |
| 70 | |
| 71 | Note that a real production scenario would require |
| 72 | logprobs to be sorted in descending order along rows, |
| 73 | something which is omitted in this function. |
| 74 | |
| 75 | Args: |
| 76 | shape: (num_tokens,num_logprobs) tuple representing |
| 77 | matrix shape |
| 78 | lower: lower range of logprob float values |
| 79 | upper: upper range of logprob float values |
| 80 | |
| 81 | Returns: |
| 82 | 2D num_tokens x num_logprobs torch Tensor of float logprob values |
| 83 | |
| 84 | """ |
| 85 | return torch.rand(*shape) * (upper - lower) + lower |
| 86 | |
| 87 | |
| 88 | def _create_random_top_token_test_vector( |
| 89 | num_logprobs: int, |
| 90 | lower: int, |
| 91 | upper: int, |
| 92 | sampled_token_id: int, |
| 93 | adjust_num_logprobs: bool = True, |
| 94 | ) -> tuple[torch.Tensor, int]: |
| 95 | """Create a random vector of top logprob token indices. |
| 96 | |
| 97 | Use to create fake sample logprobs for testing. The sampled token |
| 98 | ID must always be one of the top logprobs, which this dummy test |
| 99 | vector generator enforces. OpenAI API |
| 100 | compatible engines must be able to return an additional sample |
| 101 | logprob for the sampled token if the sampled token was not |
| 102 | among the top sample logprobs; `adjust_num_logprobs` emulates |
| 103 | this behavior by increasing the vector length by 1 if |
| 104 | `adjust_num_logprobs` is set. |
| 105 | |
| 106 | Args: |
| 107 | num_logprobs: number of top logprobs |
| 108 | lower: lower range of token ids |
| 109 | upper: upper range of token ids |
| 110 | sampled_token_id: the token actually sampled |
| 111 | adjust_num_logprobs: if True, emulate situation where sampled |
| 112 | token logprob must be injected into top |
| 113 | logprobs |
| 114 | |
| 115 | Returns: |
| 116 | 1D length-x torch Tensor of token ids where x is |
| 117 | `num_logprobs+1` if `adjust_num_logprobs` and |
| 118 | `num_logprobs` otherwise |
| 119 | sampled_token_rank: the rank of sampled_token_id in the vocab |
| 120 | vector when sorted in descending order by |
| 121 | logprob |
| 122 | |
| 123 | """ |
| 124 | # Calculate the final number of logprobs required |
| 125 | total_logprobs = num_logprobs + 1 if adjust_num_logprobs else num_logprobs |
| 126 | |
| 127 | # Generate random indices using torch |
| 128 | choice_tensor = torch.randperm(upper - lower)[:total_logprobs] + lower |
| 129 | |
| 130 | # Ensure the sampled token ID is included in the tensor |
| 131 | choice_tensor[0] = sampled_token_id |
| 132 | |
| 133 | # Check if the sampled_token_id occurs in choice_tensor[1:] |
| 134 | if sampled_token_id in choice_tensor[1:]: |
| 135 | # Use 1-based rank to match the top-k rank convention used by the |
| 136 | # consumer (logprobs.py assigns ranks starting from 1). |
| 137 | sampled_token_rank = (choice_tensor[1:] == sampled_token_id).nonzero( |
| 138 | as_tuple=True |
| 139 | )[0].item() + 1 |
| 140 | else: |
| 141 | # If not found, assign a random int between num_logprobs and 50700 |
| 142 | sampled_token_rank = random.randint(num_logprobs, 50700) |
| 143 | |
| 144 | return choice_tensor, sampled_token_rank |
| 145 | |
| 146 | |
| 147 | def _create_random_top_token_test_matrix( |
| 148 | shape: tuple[int, int], |
| 149 | lower: int, |
| 150 | upper: int, |
| 151 | tokens_list: list[int], |
| 152 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 153 | """Create a random matrix of top logprob token indices. |
| 154 | |
| 155 | Use to create fake prompt logprobs for testing. |
| 156 | |
| 157 | Token ids are generated randomly and sampled without |
| 158 | replacement. |
| 159 | |
| 160 | Args: |
| 161 | shape: (num_tokens, num_logprobs) tuple representing |
| 162 | matrix shape |
| 163 | lower: lower range of token ids |
| 164 | upper: upper range of token ids |
| 165 | |
| 166 | Returns: |
| 167 | tuple containing: |
| 168 | - 2D num_tokens x num_logprobs+1 torch Tensor of token ids |
| 169 | - 1D tensor of ranks of prompt tokens in their respective |
| 170 | rows, or random values |
| 171 | |
| 172 | """ |
| 173 | num_elements = shape[0] * shape[1] |
| 174 | choice_tensor = torch.randperm(upper - lower)[:num_elements] + lower |
| 175 | matrix = torch.cat( |
| 176 | ( |
| 177 | torch.tensor(tokens_list, dtype=torch.int).unsqueeze(-1), |
| 178 | choice_tensor.view(shape), |
| 179 | ), |
| 180 | dim=1, |
| 181 | ) |
| 182 | |
| 183 | # Initialize the tensor for storing the ranks |
| 184 | prompt_token_ranks = torch.empty(shape[0], dtype=torch.int) |
| 185 | |
| 186 | # Iterate over each row to check presence of |
| 187 | # tokens_list[rdx] and determine its index |
| 188 | for rdx in range(shape[0]): |
| 189 | row = matrix[rdx, 1:] # Skip the first column as it contains the token list |
| 190 | token_index = (row == tokens_list[rdx]).nonzero(as_tuple=True)[0] |
| 191 | if token_index.numel() > 0: |
| 192 | prompt_token_ranks[rdx] = token_index.item() + 1 |
| 193 | else: |
| 194 | prompt_token_ranks[rdx] = random.randint(shape[1], 50700) |
| 195 | |
| 196 | return matrix, prompt_token_ranks |
| 197 | |
| 198 | |
| 199 | def generate_dummy_sample_logprobs( |
| 200 | sampled_tokens_list: list, |
| 201 | num_logprobs: int, |
| 202 | tokenizer: PythonBackend, |
| 203 | ) -> list[tuple[list[int], list[float], int]]: |
| 204 | """Generate dummy sample logprobs. |
| 205 | |
| 206 | Generate a test data structure which imitates the list of sample logprobs |
| 207 | which would be assembled in the engine core during decode phase. |
| 208 | |
| 209 | Args: |
| 210 | sampled_tokens_list: list of sampled tokens |
| 211 | num_logprobs: return `num_logprobs` or `num_logprobs+1` logprobs per token |
| 212 | tokenizer: model tokenizer to use for detokenization |
| 213 | |
| 214 | Returns: |
| 215 | list of (top token ids vector, logprobs vector, sampled token rank) |
| 216 | Python lists tuples; in each tuple the logprobs and top token ids |
| 217 | vectors have the same length which is either `num_logprobs` or |
| 218 | `num_logprobs+1`. Sampled token rank is the rank (index+1) of the |
| 219 | sampled token within the vocab vector when sorted by logprob in |
| 220 | descending order. |
| 221 | |
| 222 | """ |
| 223 | res = [] |
| 224 | for sampled_token_id in sampled_tokens_list: |
| 225 | ( |
| 226 | token_vector, |
| 227 | sampled_token_rank, |
| 228 | ) = _create_random_top_token_test_vector( |
| 229 | num_logprobs, 0, len(tokenizer.vocab) - 1, sampled_token_id |
| 230 | ) |
| 231 | |
| 232 | res.append( |
| 233 | ( |
| 234 | token_vector, |
| 235 | _create_random_top_logprob_test_vector(num_logprobs + 1, -100, 0), |
| 236 | sampled_token_rank, |
| 237 | ) |
| 238 | ) |
| 239 | |
| 240 | # Convert tensors in the list tuples to Python lists |
| 241 | res_list_format = [ |
| 242 | (log_probs_tensor.tolist(), token_ids_tensor.tolist(), sampled_token_rank) |
| 243 | for log_probs_tensor, token_ids_tensor, sampled_token_rank in res |
| 244 | ] |
| 245 | |
| 246 | return res_list_format |
| 247 | |
| 248 | |
| 249 | def generate_dummy_prompt_logprobs_tensors( |
| 250 | prompt_tokens_list: list, |
| 251 | num_logprobs: int, |
| 252 | tokenizer: PythonBackend, |
| 253 | ) -> LogprobsTensors: |
| 254 | """Generate dummy prompt logprobs tensors. |
| 255 | |
| 256 | Generate a test data structure which imitates the torch Tensors of prompt |
| 257 | logprobs which would be assembled in the engine core during chunked |
| 258 | prefill. |
| 259 | |
| 260 | Args: |
| 261 | prompt_tokens_list: list of prompt tokens |
| 262 | num_logprobs: return `num_logprobs` logprobs per token |
| 263 | tokenizer: model tokenizer to use for detokenization |
| 264 | |
| 265 | Returns: |
| 266 | Single tuple of (logprobs matrix, top token ids matrix) torch Tensor, |
| 267 | where both matrices have dimensions |
| 268 | num_prompt_tokens x num_logprobs |
| 269 | |
| 270 | """ |
| 271 | # For now, assume the whole prompt is processed in one chunk; thus, |
| 272 | # the number of non-`None` prompt logprobs is `len(prompt_tokens_list)-1`. |
| 273 | # Prior to injecting `None` at the beginning of prompt logprobs (which |
| 274 | # happens later in the detokenizer, not here), the prompt logprobs in |
| 275 | # the ith position are predicting the probability distribution of the |
| 276 | # prompt token in (i+1)st position. Thus, we concat |
| 277 | # `prompt_tokens_list[1:]` to the dummy token ids, just as the engine |
| 278 | # would. |
| 279 | num_prompt_logprobs = len(prompt_tokens_list) - 1 |
| 280 | ( |
| 281 | token_vector, |
| 282 | prompt_token_ranks, |
| 283 | ) = _create_random_top_token_test_matrix( |
| 284 | (num_prompt_logprobs, num_logprobs), |
| 285 | 0, |
| 286 | len(tokenizer.vocab) - 1, |
| 287 | prompt_tokens_list[1:], |
| 288 | ) |
| 289 | return LogprobsTensors( |
| 290 | token_vector, |
| 291 | _create_random_top_logprob_test_matrix( |
| 292 | (num_prompt_logprobs, num_logprobs + 1), -100, 0 |
| 293 | ), |
| 294 | prompt_token_ranks, |
| 295 | ) |
| 296 | |
| 297 | |
| 298 | @dataclass |
| 299 | class DummyOutputProcessorTestVectors: |
| 300 | """Dummy test vectors for output processor tests.""" |
| 301 | |
| 302 | tokenizer: GeneralTokenizerType |
| 303 | vllm_config: EngineArgs |
| 304 | full_tokens: list[list[int]] # Prompt + generated tokens |
| 305 | prompt_tokens: list[list[int]] |
| 306 | generation_tokens: list[list[int]] |
| 307 | # Each request is associated with a tuple of |
| 308 | # (top tokens, top logprobs, ranks) prompt logprobs tensors |
| 309 | prompt_logprobs: list[LogprobsTensors] |
| 310 | # Each request is associated with a sample logprobs; a request's |
| 311 | # sample logprobs are a list of (top tokens, top logprobs, ranks) |
| 312 | # sample logprobs tensors at each sequence position |
| 313 | generation_logprobs: list[list[tuple[list[int], list[float], int]]] |
| 314 | prompt_strings: list[str] |
| 315 | prompt_strings_len: list[int] |
| 316 | generation_strings: list[str] |
| 317 | |
| 318 | |
| 319 | class MockEngineCore: |
| 320 | """Mock engine core outputs form premade tokens lists.""" |
| 321 | |
| 322 | def __init__( |
| 323 | self, |
| 324 | tokens_list: list[list[int]], |
| 325 | prompts_list: list[list[int]], |
| 326 | # For each request, for each sampled token offset, |
| 327 | # a tuple of |
| 328 | # (list of topk token ids, list of sample logprob vals, rank) |
| 329 | generated_logprobs_raw: list[list[tuple[list[int], list[float], int]]] |
| 330 | | None = None, |
| 331 | # For each request, a tuple of |
| 332 | # (prompt logprob val matrix, prompt logprob tok id matrix); |
| 333 | # each matrix has dimensions |
| 334 | # (num prompt toks) x (num prompt logprobs+1) |
| 335 | prompt_logprobs_raw: list[LogprobsTensors] | None = None, |
| 336 | eos_token_id: int | None = None, |
| 337 | stop_token_ids: list[int] | None = None, |
| 338 | request_ids: list[str] | None = None, |
| 339 | ) -> None: |
| 340 | self.num_requests = len(tokens_list) |
| 341 | self.tokens_list = tokens_list |
| 342 | self.prompts_list = prompts_list |
| 343 | self.generated_logprobs_raw = generated_logprobs_raw |
| 344 | self.do_logprobs = generated_logprobs_raw is not None |
| 345 | self.prompt_logprobs_raw = prompt_logprobs_raw |
| 346 | self.do_prompt_logprobs = prompt_logprobs_raw is not None |
| 347 | self.request_finished = [False for _ in range(self.num_requests)] |
| 348 | self.request_token_idx = [0 for _ in range(self.num_requests)] |
| 349 | self.eos_token_id = eos_token_id |
| 350 | self.stop_token_ids = stop_token_ids |
| 351 | self.request_ids = ( |
| 352 | request_ids |
| 353 | if request_ids is not None |
| 354 | else [f"request-{i}" for i in range(self.num_requests)] |
| 355 | ) |
| 356 | |
| 357 | def get_outputs(self, num_active: int = -1) -> list[EngineCoreOutput]: |
| 358 | do_logprobs = self.do_logprobs |
| 359 | do_prompt_logprobs = self.do_prompt_logprobs |
| 360 | |
| 361 | outputs = [] |
| 362 | for req_idx, (token_ids, prompt_token_ids) in enumerate( |
| 363 | zip(self.tokens_list, self.prompts_list) |
| 364 | ): |
| 365 | if num_active != -1 and req_idx >= num_active: |
| 366 | break |
| 367 | if not self.request_finished[req_idx]: |
| 368 | token_idx = self.request_token_idx[req_idx] |
| 369 | if do_logprobs: |
| 370 | assert self.generated_logprobs_raw is not None |
| 371 | (logprobs_token_ids_, logprobs_, sampled_token_ranks_) = ( |
| 372 | self.generated_logprobs_raw[req_idx][token_idx] |
| 373 | ) |
| 374 | logprobs = LogprobsLists( |
| 375 | np.array([logprobs_token_ids_]), |
| 376 | np.array([logprobs_]), |
| 377 | np.array([sampled_token_ranks_]), |
| 378 | ) |
| 379 | else: |
| 380 | logprobs = None |
| 381 | if do_prompt_logprobs: |
| 382 | if token_idx == 0: |
| 383 | assert self.prompt_logprobs_raw is not None |
| 384 | prompt_logprobs = self.prompt_logprobs_raw[req_idx] |
| 385 | else: |
| 386 | prompt_logprobs = None |
| 387 | else: |
| 388 | prompt_logprobs = None |
| 389 | |
| 390 | # Add prefill_stats on first output (prefill) for this request |
| 391 | if token_idx == 0: |
| 392 | prefill_stats = PrefillStats() |
| 393 | prefill_stats.set( |
| 394 | num_prompt_tokens=len(prompt_token_ids), |
| 395 | num_local_cached_tokens=0, |
| 396 | num_external_cached_tokens=0, |
| 397 | ) |
| 398 | else: |
| 399 | prefill_stats = None |
| 400 | |
| 401 | new_token_id = token_ids[token_idx] |
| 402 | output = EngineCoreOutput( |
| 403 | request_id=self.request_ids[req_idx], |
| 404 | new_token_ids=[new_token_id], |
| 405 | new_logprobs=logprobs, |
| 406 | new_prompt_logprobs_tensors=prompt_logprobs, |
| 407 | prefill_stats=prefill_stats, |
| 408 | ) |
| 409 | if token_idx == len(token_ids) - 1: |
| 410 | output.finish_reason = FinishReason.LENGTH |
| 411 | self.request_finished[req_idx] = True |
| 412 | if new_token_id == self.eos_token_id: |
| 413 | output.finish_reason = FinishReason.STOP |
| 414 | self.request_finished[req_idx] = True |
| 415 | if new_token_id in (self.stop_token_ids or ()): |
| 416 | output.finish_reason = FinishReason.STOP |
| 417 | output.stop_reason = new_token_id |
| 418 | self.request_finished[req_idx] = True |
| 419 | outputs.append(output) |
| 420 | |
| 421 | self.request_token_idx[req_idx] += 1 |
| 422 | |
| 423 | return outputs |