| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import itertools |
| 5 | import math |
| 6 | from collections.abc import Generator |
| 7 | from types import SimpleNamespace |
| 8 | from typing import get_args |
| 9 | from unittest.mock import patch |
| 10 | |
| 11 | import numpy as np |
| 12 | import pytest |
| 13 | import torch |
| 14 | |
| 15 | from tests.utils import large_gpu_mark |
| 16 | from tests.v1.sample.utils import ( |
| 17 | BatchLogprobsComposition, |
| 18 | BatchLogprobsSpecType, |
| 19 | assert_incr_detok_str_matches_non_incr_detok_str, |
| 20 | compute_correct_cumulative_logprob, |
| 21 | get_test_batch, |
| 22 | ) |
| 23 | from vllm import SamplingParams |
| 24 | from vllm.config.model import LogprobsMode |
| 25 | from vllm.distributed import cleanup_dist_env_and_memory |
| 26 | from vllm.exceptions import VLLMValidationError |
| 27 | from vllm.logprobs import Logprob |
| 28 | from vllm.platforms import current_platform |
| 29 | from vllm.v1.engine.input_processor import InputProcessor |
| 30 | |
| 31 | from ...conftest import HfRunner, VllmRunner |
| 32 | |
| 33 | MODEL = "meta-llama/Llama-3.2-1B-Instruct" |
| 34 | DTYPE = "half" |
| 35 | |
| 36 | NONE = BatchLogprobsComposition.NONE |
| 37 | SAMPLE = BatchLogprobsComposition.SAMPLE |
| 38 | PROMPT = BatchLogprobsComposition.PROMPT |
| 39 | SAMPLE_PROMPT = BatchLogprobsComposition.SAMPLE_PROMPT |
| 40 | |
| 41 | # On ROCm, floating-point reductions in attention and GEMM kernels are |
| 42 | # non-associative and sensitive to batch geometry. If the ref LLM and |
| 43 | # spec-decode LLM use different scheduling or batch geometry, they can |
| 44 | # follow different reduction orders and produce numerically divergent |
| 45 | # logprobs that get misattributed to spec-decode incorrectness. |
| 46 | # |
| 47 | # Force LLM instances into an identical, deterministic execution |
| 48 | # mode so the test isolates spec-decode correctness only: |
| 49 | if current_platform.is_rocm(): |
| 50 | GPU_DETERMINISM_KWARGS: dict = dict(max_num_seqs=1, attention_backend="TRITON_ATTN") |
| 51 | elif current_platform.is_xpu(): |
| 52 | GPU_DETERMINISM_KWARGS = dict(max_num_seqs=1, attention_backend="FLASH_ATTN") |
| 53 | else: |
| 54 | GPU_DETERMINISM_KWARGS = {} |
| 55 | |
| 56 | |
| 57 | @pytest.fixture( |
| 58 | scope="module", |
| 59 | # Parameterize APC |
| 60 | params=[False, True], |
| 61 | ) |
| 62 | def vllm_model(vllm_runner, request) -> Generator[VllmRunner, None, None]: |
| 63 | with vllm_runner( |
| 64 | MODEL, |
| 65 | dtype=DTYPE, |
| 66 | max_logprobs=7, |
| 67 | # Very small number of batched tokens to ensure |
| 68 | # that we test chunking. |
| 69 | max_num_batched_tokens=16, |
| 70 | max_num_seqs=16, |
| 71 | max_model_len=128, |
| 72 | enable_chunked_prefill=True, |
| 73 | enforce_eager=True, |
| 74 | # TODO: enable this once we support it for |
| 75 | # prompt logprobs. |
| 76 | enable_prefix_caching=request.param, |
| 77 | gpu_memory_utilization=0.4, |
| 78 | ) as vllm_model: |
| 79 | yield vllm_model |
| 80 | |
| 81 | |
| 82 | @pytest.fixture(scope="module") |
| 83 | def hf_model(hf_runner) -> Generator[HfRunner, None, None]: |
| 84 | with hf_runner(MODEL, dtype=DTYPE) as hf_model: |
| 85 | yield hf_model |
| 86 | |
| 87 | |
| 88 | def _model_config(vocab_size: int = 10, max_logprobs: int = 20): |
| 89 | return SimpleNamespace( |
| 90 | max_logprobs=max_logprobs, |
| 91 | logits_processors=None, |
| 92 | is_diffusion=False, |
| 93 | get_vocab_size=lambda: vocab_size, |
| 94 | ) |
| 95 | |
| 96 | |
| 97 | def _repeat_logprob_config( |
| 98 | test_prompts, |
| 99 | logprob_prompt_logprob_list: BatchLogprobsSpecType, |
| 100 | ) -> BatchLogprobsSpecType: |
| 101 | """Ensure each test prompt has a logprob config. |
| 102 | |
| 103 | A logprob config specifies the optional (i.e. |
| 104 | may-be-`None`) number of sample logprobs and |
| 105 | the optional number of prompt logprobs. |
| 106 | |
| 107 | If more test prompts than logprob configs are |
| 108 | provided, the provided logprob configs are |
| 109 | tiled to match the number of test prompts. |
| 110 | |
| 111 | If fewer test prompts than logprob configs |
| 112 | are provided, the list of logprob configs |
| 113 | is truncated to match the number of test |
| 114 | prompts. |
| 115 | |
| 116 | Otherwise, the list of logprob configs |
| 117 | is returned as-is. |
| 118 | |
| 119 | Args: |
| 120 | test_prompts: list of prompts under test |
| 121 | logprob_prompt_logprob_list: list of |
| 122 | (optional num sample logprob, |
| 123 | optional num prompt logprob) |
| 124 | tuples |
| 125 | |
| 126 | Returns: |
| 127 | list of |
| 128 | (optional num sample logprob,optional num prompt logprob) |
| 129 | tuples which is either identical to |
| 130 | `logprob_prompt_logprob_list`, or else repeats |
| 131 | `logprob_prompt_logprob_list` enough times to match the |
| 132 | number of `test_prompts`, or else is truncated to match |
| 133 | the number of `test_prompts` |
| 134 | |
| 135 | """ |
| 136 | num_test_prompts = len(test_prompts) |
| 137 | # Make sure there is a logprobs configuration for each test prompt |
| 138 | logprob_prompt_logprob_list = list( |
| 139 | itertools.islice(itertools.cycle(logprob_prompt_logprob_list), num_test_prompts) |
| 140 | ) |
| 141 | # Now the number of prompts should match the number of sample params combos |
| 142 | assert num_test_prompts == len(logprob_prompt_logprob_list) |
| 143 | return logprob_prompt_logprob_list |
| 144 | |
| 145 | |
| 146 | def _run_and_validate( |
| 147 | vllm_model: VllmRunner, |
| 148 | test_prompts: list[str], |
| 149 | vllm_sampling_params: SamplingParams | list[SamplingParams], |
| 150 | hf_logprobs: list[list[torch.Tensor]], |
| 151 | hf_outputs: list[tuple[list[int], str]], |
| 152 | logprob_prompt_logprob_list: BatchLogprobsSpecType, |
| 153 | temperature: float, |
| 154 | max_tokens: int, |
| 155 | do_apc: bool, |
| 156 | ) -> None: |
| 157 | vllm_results = vllm_model.llm.generate( |
| 158 | test_prompts, sampling_params=vllm_sampling_params |
| 159 | ) |
| 160 | |
| 161 | for vllm_result, hf_logprob, hf_output, logprob_prompt_logprob in zip( |
| 162 | vllm_results, hf_logprobs, hf_outputs, logprob_prompt_logprob_list |
| 163 | ): |
| 164 | # Extract request-level (prompt)logprobs config |
| 165 | num_top_logprobs, num_top_prompt_logprobs = logprob_prompt_logprob |
| 166 | |
| 167 | assert vllm_result.prompt_token_ids is not None |
| 168 | # Test whether sampled token output is consistent between vLLM and HF |
| 169 | # vLLM prompt+completion should match HF output |
| 170 | if temperature == 0.0: |
| 171 | assert ( |
| 172 | vllm_result.prompt_token_ids + list(vllm_result.outputs[0].token_ids) |
| 173 | == hf_output[0] |
| 174 | ) |
| 175 | else: |
| 176 | # Sampled tokens won't match if not greedy |
| 177 | assert ( |
| 178 | vllm_result.prompt_token_ids |
| 179 | == hf_output[0][: len(vllm_result.prompt_token_ids)] |
| 180 | ) |
| 181 | |
| 182 | # Validate sample logprobs |
| 183 | if num_top_logprobs is not None: |
| 184 | assert num_top_logprobs is not None |
| 185 | # Confirm that the structure of the sample logprobs in the result is |
| 186 | # correct |
| 187 | assert vllm_result.outputs[0].logprobs is not None |
| 188 | assert len(vllm_result.outputs[0].logprobs) == max_tokens |
| 189 | for logprobs, token_id in zip( |
| 190 | vllm_result.outputs[0].logprobs, vllm_result.outputs[0].token_ids |
| 191 | ): |
| 192 | assert logprobs is not None |
| 193 | |
| 194 | # Confirm that the output token appears among the logprobs |
| 195 | assert token_id in logprobs |
| 196 | rank = logprobs[token_id].rank |
| 197 | assert rank is not None |
| 198 | token_in_topk = rank <= num_top_logprobs |
| 199 | |
| 200 | # If the output token is not included in the top K |
| 201 | # logprob, it can return 1 more data |
| 202 | if token_in_topk and num_top_logprobs != 0: |
| 203 | assert len(logprobs) == num_top_logprobs |
| 204 | else: |
| 205 | assert len(logprobs) == num_top_logprobs + 1 |
| 206 | |
| 207 | if num_top_logprobs > 0: |
| 208 | # We should have an entry for each of the topk ranks |
| 209 | all_ranks = {lp.rank for lp in logprobs.values()} |
| 210 | assert all(r in all_ranks for r in range(1, num_top_logprobs + 1)) |
| 211 | |
| 212 | output_text = vllm_result.outputs[0].text |
| 213 | output_string_from_most_likely_tokens_lst: list[str] = [] |
| 214 | for output_logprobs in vllm_result.outputs[0].logprobs: |
| 215 | assert output_logprobs is not None |
| 216 | top_logprob = next(iter(output_logprobs.values())) |
| 217 | assert top_logprob.decoded_token is not None |
| 218 | output_string_from_most_likely_tokens_lst.append( |
| 219 | top_logprob.decoded_token |
| 220 | ) |
| 221 | |
| 222 | output_string_from_most_likely_tokens = "".join( |
| 223 | output_string_from_most_likely_tokens_lst |
| 224 | ) |
| 225 | assert_incr_detok_str_matches_non_incr_detok_str( |
| 226 | output_text, |
| 227 | output_string_from_most_likely_tokens, |
| 228 | "The output text from the top logprob for each token " |
| 229 | "position should be the same as the output text in the " |
| 230 | "result.", |
| 231 | ) |
| 232 | |
| 233 | # Compare vLLM sample logprobs to HF |
| 234 | vllm_sample_logprobs = vllm_result.outputs[0].logprobs |
| 235 | for i, sample_logprobs in enumerate(vllm_sample_logprobs): |
| 236 | assert sample_logprobs is not None |
| 237 | for token_id, sample_logprob in sample_logprobs.items(): |
| 238 | if temperature == 0.0 or i == 0: |
| 239 | logprob = sample_logprob.logprob |
| 240 | torch.testing.assert_close( |
| 241 | logprob, |
| 242 | hf_logprob[i][-1][token_id].item(), |
| 243 | atol=1e-2, |
| 244 | rtol=1e-2, |
| 245 | ) |
| 246 | assert isinstance(sample_logprob.decoded_token, str), ( |
| 247 | "The token should be decoded by the time it is" |
| 248 | " returned to the user." |
| 249 | ) |
| 250 | |
| 251 | # At this point we know the sample logprobs are correct for this |
| 252 | # request. Validate that cumulative_logprob is actually the sum. |
| 253 | # For each request, assert that the returned cumulative logprob |
| 254 | # matches the correct value, which is computed below. |
| 255 | torch.testing.assert_close( |
| 256 | vllm_result.outputs[0].cumulative_logprob, |
| 257 | compute_correct_cumulative_logprob(vllm_result.outputs[0]), |
| 258 | atol=1e-6, |
| 259 | rtol=1e-6, |
| 260 | ) |
| 261 | else: |
| 262 | # Logprobs disabled for this request; should be None |
| 263 | assert vllm_result.outputs[0].logprobs is None |
| 264 | |
| 265 | # Validate prompt logprobs |
| 266 | if num_top_prompt_logprobs is not None: |
| 267 | # Confirm that structure of prompt logprobs in result is correct |
| 268 | assert vllm_result.prompt_logprobs is not None |
| 269 | # - The first prompt logprob is always None |
| 270 | assert vllm_result.prompt_logprobs[0] is None |
| 271 | # - Prompt logprobs are returned for all indices in |
| 272 | # the prompt |
| 273 | assert len(vllm_result.prompt_logprobs) == len(vllm_result.prompt_token_ids) |
| 274 | for prompt_logprobs, prompt_token_id in zip( |
| 275 | vllm_result.prompt_logprobs[1:], vllm_result.prompt_token_ids[1:] |
| 276 | ): |
| 277 | assert prompt_logprobs is not None |
| 278 | |
| 279 | # Confirm that the prompt token appears among the logprobs |
| 280 | assert prompt_token_id in prompt_logprobs |
| 281 | prompt_rank = prompt_logprobs[prompt_token_id].rank |
| 282 | assert prompt_rank is not None |
| 283 | token_in_topk = prompt_rank <= num_top_prompt_logprobs |
| 284 | |
| 285 | # If the prompt token is not included in the top K |
| 286 | # logprob, it can return 1 more data |
| 287 | if token_in_topk and num_top_prompt_logprobs != 0: |
| 288 | assert len(prompt_logprobs) == num_top_prompt_logprobs |
| 289 | else: |
| 290 | assert len(prompt_logprobs) == num_top_prompt_logprobs + 1 |
| 291 | |
| 292 | if num_top_prompt_logprobs > 0: |
| 293 | # We should have an entry for each of the topk ranks |
| 294 | all_ranks = {lp.rank for lp in prompt_logprobs.values()} |
| 295 | assert all( |
| 296 | r in all_ranks for r in range(1, num_top_prompt_logprobs + 1) |
| 297 | ) |
| 298 | |
| 299 | # Compare prompt logprobs to HF |
| 300 | # The first prompt logprob is always None, so we compare it from |
| 301 | # 1:. |
| 302 | vllm_prompt_logprobs = vllm_result.prompt_logprobs[1:] |
| 303 | for i, vllm_prompt_logprob_dict in enumerate(vllm_prompt_logprobs): |
| 304 | assert vllm_prompt_logprob_dict is not None |
| 305 | for token_id, prompt_lp in vllm_prompt_logprob_dict.items(): |
| 306 | torch.testing.assert_close( |
| 307 | prompt_lp.logprob, |
| 308 | hf_logprob[0][i][token_id].item(), |
| 309 | atol=2e-2, |
| 310 | rtol=2e-2, |
| 311 | ) |
| 312 | else: |
| 313 | assert vllm_result.prompt_logprobs is None |
| 314 | |
| 315 | |
| 316 | @pytest.mark.parametrize( |
| 317 | "batch_logprobs_composition", [NONE, SAMPLE, PROMPT, SAMPLE_PROMPT] |
| 318 | ) |
| 319 | @pytest.mark.parametrize("temperature", [0.0, 2.0]) |
| 320 | def test_get_logprobs_and_prompt_logprobs( |
| 321 | hf_model, |
| 322 | vllm_model, |
| 323 | batch_logprobs_composition: BatchLogprobsComposition, |
| 324 | temperature: float, |
| 325 | example_prompts: list[str], |
| 326 | ) -> None: |
| 327 | """Test V1 Engine logprobs & prompt logprobs. |
| 328 | |
| 329 | Exercise a variety of combinations of `logprobs` and `prompt_logprobs` |
| 330 | settings and validate that |
| 331 | * The generated logprobs and prompt logprobs are consistent with the |
| 332 | configuration settings, in terms of whether or not the logprobs |
| 333 | (of either type) were requested and how many were requested |
| 334 | * The generated logprobs are consistent with the generated tokens |
| 335 | * The generated (prompt)logprobs are consistent with HuggingFace |
| 336 | (prompt)logprobs, as a reference |
| 337 | |
| 338 | batch_logprobs_composition controls the logprobs configurations for |
| 339 | requests in the batch under test. |
| 340 | |
| 341 | APC tests run two test iterations so that cache hits occur. |
| 342 | |
| 343 | To save time, only test one APC-enabled scenario |
| 344 | (sample & prompt logprobs enabled, temperature>0.0). |
| 345 | |
| 346 | Args: |
| 347 | hf_model: HuggingFace reference model fixture |
| 348 | vllm_model: vLLM model fixture |
| 349 | batch_logprobs_composition: logprobs configuration for test batch |
| 350 | temperature: "temperature" sampling parameter |
| 351 | example_prompts: example prompt fixture |
| 352 | |
| 353 | """ |
| 354 | vllm_config = vllm_model.llm.llm_engine.vllm_config |
| 355 | do_apc = vllm_config.cache_config.enable_prefix_caching |
| 356 | if do_apc and (temperature < 2.0 or batch_logprobs_composition != SAMPLE_PROMPT): |
| 357 | # Skip some test-cases to save time. |
| 358 | pytest.skip() |
| 359 | test_prompts = example_prompts |
| 360 | |
| 361 | max_tokens = 5 |
| 362 | hf_outputs = hf_model.generate_greedy( |
| 363 | test_prompts, |
| 364 | max_tokens=max_tokens, |
| 365 | ) |
| 366 | hf_logprobs = hf_model.generate_greedy_logprobs( |
| 367 | test_prompts, |
| 368 | max_tokens=max_tokens, |
| 369 | ) |
| 370 | |
| 371 | # Batch has mixed sample params |
| 372 | # (different logprobs/prompt logprobs combos) |
| 373 | logprob_prompt_logprob_list = get_test_batch(batch_logprobs_composition) |
| 374 | |
| 375 | # Ensure that each test prompt has a logprob config for testing |
| 376 | logprob_prompt_logprob_list = _repeat_logprob_config( |
| 377 | test_prompts, logprob_prompt_logprob_list |
| 378 | ) |
| 379 | # Generate SamplingParams |
| 380 | vllm_sampling_params = [ |
| 381 | SamplingParams( |
| 382 | max_tokens=max_tokens, |
| 383 | logprobs=num_lp, |
| 384 | prompt_logprobs=num_plp, |
| 385 | temperature=temperature, |
| 386 | seed=1984, |
| 387 | ) |
| 388 | for num_lp, num_plp in logprob_prompt_logprob_list |
| 389 | ] |
| 390 | for _ in range(2 if do_apc else 1): |
| 391 | _run_and_validate( |
| 392 | vllm_model=vllm_model, |
| 393 | test_prompts=test_prompts, |
| 394 | vllm_sampling_params=vllm_sampling_params, |
| 395 | hf_logprobs=hf_logprobs, |
| 396 | hf_outputs=hf_outputs, |
| 397 | logprob_prompt_logprob_list=logprob_prompt_logprob_list, |
| 398 | temperature=temperature, |
| 399 | max_tokens=max_tokens, |
| 400 | do_apc=do_apc, |
| 401 | ) |
| 402 | |
| 403 | |
| 404 | def test_max_logprobs(): |
| 405 | """VLLM v1 engine should fail a request with `logprobs > max_logprobs` |
| 406 | Should also fail for `prompt_logprobs > max_logprobs` |
| 407 | APC should not matter as this test checks basic request validation. |
| 408 | """ |
| 409 | with VllmRunner( |
| 410 | "facebook/opt-125m", |
| 411 | max_logprobs=1, |
| 412 | enable_prefix_caching=False, |
| 413 | gpu_memory_utilization=0.15, |
| 414 | max_model_len=256, |
| 415 | ) as runner: |
| 416 | vllm_sampling_params = SamplingParams(logprobs=1) |
| 417 | # should pass |
| 418 | runner.generate(["Hello world"], sampling_params=vllm_sampling_params) |
| 419 | |
| 420 | bad_sampling_params = SamplingParams(logprobs=2) |
| 421 | with pytest.raises(VLLMValidationError): |
| 422 | runner.generate(["Hello world"], sampling_params=bad_sampling_params) |
| 423 | |
| 424 | |
| 425 | @pytest.mark.parametrize("token_ids", [[0], [0, 9]]) |
| 426 | def test_logprob_token_ids_validate_vocab_bounds_valid(token_ids: list[int]): |
| 427 | SamplingParams(logprob_token_ids=token_ids).verify( |
| 428 | _model_config(), |
| 429 | speculative_config=None, |
| 430 | structured_outputs_config=None, |
| 431 | tokenizer=None, |
| 432 | ) |
| 433 | |
| 434 | |
| 435 | @pytest.mark.parametrize("token_ids", [[-1], [10], [-35, 1873042417]]) |
| 436 | def test_logprob_token_ids_validate_vocab_bounds_invalid(token_ids: list[int]): |
| 437 | with pytest.raises(VLLMValidationError, match="logprob_token_ids"): |
| 438 | SamplingParams(logprob_token_ids=token_ids).verify( |
| 439 | _model_config(), |
| 440 | speculative_config=None, |
| 441 | structured_outputs_config=None, |
| 442 | tokenizer=None, |
| 443 | ) |
| 444 | |
| 445 | |
| 446 | def test_prompt_logprob_token_ids_bounded_by_max_logprobs(): |
| 447 | """The candidate count is bounded by max_logprobs.""" |
| 448 | model_config = _model_config(vocab_size=100, max_logprobs=4) |
| 449 | |
| 450 | def verify(**kwargs): |
| 451 | SamplingParams(**kwargs).verify( |
| 452 | model_config, |
| 453 | speculative_config=None, |
| 454 | structured_outputs_config=None, |
| 455 | tokenizer=None, |
| 456 | ) |
| 457 | |
| 458 | verify(prompt_logprob_token_ids=[1, 2, 3, 4]) |
| 459 | |
| 460 | # The error names the value to raise max_logprobs to. |
| 461 | with pytest.raises(VLLMValidationError, match=r"max_logprobs.*at least 5"): |
| 462 | verify(prompt_logprob_token_ids=[1, 2, 3, 4, 5]) |
| 463 | |
| 464 | # prompt_logprob_start alone is a caller mistake, not a silent no-op. |
| 465 | with pytest.raises(VLLMValidationError, match="requires prompt_logprob_token_ids"): |
| 466 | verify(prompt_logprob_start=3) |
| 467 | |
| 468 | |
| 469 | def test_prompt_logprob_token_ids_require_v2_model_runner(): |
| 470 | """Only the V2 runner scores them; the V1 runner would return None, and |
| 471 | kv-sharing fast prefill would score rows the cross-decoder never ran.""" |
| 472 | processor = SimpleNamespace( |
| 473 | model_config=SimpleNamespace( |
| 474 | return_sampling_mask=False, enable_trace_replay=False, is_diffusion=False |
| 475 | ), |
| 476 | vllm_config=SimpleNamespace( |
| 477 | reasoning_config=None, |
| 478 | use_v2_model_runner=False, |
| 479 | cache_config=SimpleNamespace(kv_sharing_fast_prefill=False), |
| 480 | ), |
| 481 | speculative_config=None, |
| 482 | structured_outputs_config=None, |
| 483 | tokenizer=None, |
| 484 | validate_logits_processors_params=lambda params: None, |
| 485 | ) |
| 486 | params = SamplingParams(prompt_logprob_token_ids=[1, 2]) |
| 487 | with patch.object(SamplingParams, "verify"): |
| 488 | with pytest.raises(VLLMValidationError, match="V2 model runner"): |
| 489 | InputProcessor._validate_params(processor, params, ("generate",)) |
| 490 | processor.vllm_config.use_v2_model_runner = True |
| 491 | InputProcessor._validate_params(processor, params, ("generate",)) |
| 492 | processor.vllm_config.cache_config.kv_sharing_fast_prefill = True |
| 493 | with pytest.raises(VLLMValidationError, match="fast-prefill"): |
| 494 | InputProcessor._validate_params(processor, params, ("generate",)) |
| 495 | |
| 496 | |
| 497 | def test_none_logprobs(vllm_model, example_prompts): |
| 498 | """Engine should return `logprobs` and `prompt_logprobs` as `None`. |
| 499 | |
| 500 | Args: |
| 501 | vllm_model: vLLM model fixture |
| 502 | example_prompts: list of example prompts (test fixture) |
| 503 | |
| 504 | """ |
| 505 | max_tokens = 5 |
| 506 | |
| 507 | sampling_params_logprobs_none = SamplingParams( |
| 508 | max_tokens=max_tokens, |
| 509 | logprobs=None, |
| 510 | prompt_logprobs=None, |
| 511 | temperature=0.0, |
| 512 | ) |
| 513 | results_logprobs_none = vllm_model.llm.generate( |
| 514 | example_prompts, |
| 515 | sampling_params=sampling_params_logprobs_none, |
| 516 | ) |
| 517 | |
| 518 | for i in range(len(results_logprobs_none)): |
| 519 | # Check sample logprobs are None |
| 520 | assert results_logprobs_none[i].outputs[0].logprobs is None |
| 521 | assert results_logprobs_none[i].outputs[0].cumulative_logprob is None |
| 522 | # Check prompt logprobs are None |
| 523 | assert results_logprobs_none[i].prompt_logprobs is None |
| 524 | |
| 525 | |
| 526 | def test_zero_logprobs(vllm_model, example_prompts): |
| 527 | """Engine should return sampled token and prompt token logprobs. |
| 528 | |
| 529 | Args: |
| 530 | vllm_model: vLLM model fixture |
| 531 | example_prompts: list of example prompts (test fixture) |
| 532 | |
| 533 | """ |
| 534 | max_tokens = 5 |
| 535 | |
| 536 | sampling_params_logprobs_zero = SamplingParams( |
| 537 | max_tokens=max_tokens, logprobs=0, prompt_logprobs=0, temperature=0.0 |
| 538 | ) |
| 539 | results_logprobs_zero = vllm_model.llm.generate( |
| 540 | example_prompts, sampling_params=sampling_params_logprobs_zero |
| 541 | ) |
| 542 | |
| 543 | for i in range(len(results_logprobs_zero)): |
| 544 | # Check that there is one sample logprob dict for each |
| 545 | # sample token |
| 546 | logprobs = results_logprobs_zero[i].outputs[0].logprobs |
| 547 | prompt_logprobs = results_logprobs_zero[i].prompt_logprobs |
| 548 | sampled_token_ids = results_logprobs_zero[i].outputs[0].token_ids |
| 549 | prompt_token_ids = results_logprobs_zero[i].prompt_token_ids |
| 550 | assert logprobs is not None |
| 551 | assert len(sampled_token_ids) == len(logprobs) |
| 552 | assert results_logprobs_zero[i].outputs[0].cumulative_logprob is not None |
| 553 | # Check that there is one prompt logprob dict for each |
| 554 | # prompt token |
| 555 | assert prompt_logprobs is not None |
| 556 | assert len(prompt_token_ids) == len(prompt_logprobs) |
| 557 | |
| 558 | |
| 559 | def test_all_logprobs(example_prompts): |
| 560 | """Engine should return all vocabulary logprobs and prompt logprobs. |
| 561 | |
| 562 | Args: |
| 563 | example_prompts: list of example prompts (test fixture) |
| 564 | |
| 565 | """ |
| 566 | with VllmRunner( |
| 567 | "facebook/opt-125m", |
| 568 | max_logprobs=-1, |
| 569 | enable_prefix_caching=False, |
| 570 | gpu_memory_utilization=0.15, |
| 571 | max_model_len=256, |
| 572 | ) as runner: |
| 573 | sampling_params_logprobs_all = SamplingParams( |
| 574 | max_tokens=5, logprobs=-1, prompt_logprobs=-1 |
| 575 | ) |
| 576 | results_logprobs_all = runner.llm.generate( |
| 577 | example_prompts, sampling_params=sampling_params_logprobs_all |
| 578 | ) |
| 579 | vocab_size = runner.llm.llm_engine.model_config.get_vocab_size() |
| 580 | |
| 581 | for i in range(len(results_logprobs_all)): |
| 582 | logprobs = results_logprobs_all[i].outputs[0].logprobs |
| 583 | prompt_logprobs = results_logprobs_all[i].prompt_logprobs |
| 584 | assert logprobs is not None |
| 585 | for logprob in logprobs: |
| 586 | assert len(logprob) == vocab_size |
| 587 | assert prompt_logprobs is not None |
| 588 | assert prompt_logprobs[0] is None |
| 589 | for prompt_logprob in prompt_logprobs[1:]: |
| 590 | assert len(prompt_logprob) == vocab_size |
| 591 | |
| 592 | |
| 593 | @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode)) |
| 594 | def test_logprobs_mode(logprobs_mode: LogprobsMode): |
| 595 | """Test with LLM engine with different logprobs_mode. |
| 596 | For logprobs, we should have non-positive values. |
| 597 | For logits, we should expect at least one positive values. |
| 598 | """ |
| 599 | from vllm import LLM |
| 600 | |
| 601 | llm = LLM( |
| 602 | "facebook/opt-125m", |
| 603 | max_logprobs=5, |
| 604 | enable_prefix_caching=False, |
| 605 | # 2 other llms alive during whole session; must also cover the |
| 606 | # cudagraph memory reservation from startup profiling. |
| 607 | gpu_memory_utilization=0.1, |
| 608 | max_model_len=16, |
| 609 | logprobs_mode=logprobs_mode, |
| 610 | ) |
| 611 | try: |
| 612 | vllm_sampling_params = SamplingParams(logprobs=1) |
| 613 | results = llm.generate(["Hello world"], sampling_params=vllm_sampling_params) |
| 614 | |
| 615 | total_token_with_logprobs = 0 |
| 616 | positive_values = 0 |
| 617 | for output in results[0].outputs: |
| 618 | assert output.logprobs is not None |
| 619 | for logprobs in output.logprobs: |
| 620 | assert logprobs is not None |
| 621 | for token_id in logprobs: |
| 622 | logprob = logprobs[token_id] |
| 623 | if logprobs_mode in ("raw_logprobs", "processed_logprobs"): |
| 624 | assert logprob.logprob <= 0 |
| 625 | if logprob.logprob > 0: |
| 626 | positive_values = positive_values + 1 |
| 627 | total_token_with_logprobs = total_token_with_logprobs + 1 |
| 628 | assert total_token_with_logprobs >= len(results[0].outputs) |
| 629 | if logprobs_mode in ("raw_logits", "processed_logits"): |
| 630 | assert positive_values > 0 |
| 631 | finally: |
| 632 | del llm |
| 633 | torch.accelerator.empty_cache() |
| 634 | cleanup_dist_env_and_memory() |
| 635 | |
| 636 | |
| 637 | def test_prompt_logprobs_mode(): |
| 638 | """prompt_logprobs must respect logprobs_mode: *_logits and *_logprobs |
| 639 | must return different values. Prompt tokens skip sampling processors, |
| 640 | so processed_* == raw_* on the prompt side.""" |
| 641 | from vllm import LLM |
| 642 | |
| 643 | values: dict[str, float] = {} |
| 644 | for mode in get_args(LogprobsMode): |
| 645 | llm = LLM( |
| 646 | "facebook/opt-125m", |
| 647 | enable_prefix_caching=False, |
| 648 | gpu_memory_utilization=0.1, |
| 649 | max_model_len=16, |
| 650 | logprobs_mode=mode, |
| 651 | ) |
| 652 | try: |
| 653 | results = llm.generate( |
| 654 | ["Hello world"], |
| 655 | sampling_params=SamplingParams( |
| 656 | max_tokens=1, prompt_logprobs=0, temperature=0 |
| 657 | ), |
| 658 | ) |
| 659 | assert results[0].prompt_logprobs is not None |
| 660 | assert results[0].prompt_logprobs[1] is not None |
| 661 | assert results[0].prompt_token_ids is not None |
| 662 | tok_id = results[0].prompt_token_ids[1] |
| 663 | values[mode] = results[0].prompt_logprobs[1][tok_id].logprob |
| 664 | finally: |
| 665 | del llm |
| 666 | torch.accelerator.empty_cache() |
| 667 | cleanup_dist_env_and_memory() |
| 668 | |
| 669 | assert values["raw_logprobs"] <= 0 |
| 670 | assert values["processed_logprobs"] <= 0 |
| 671 | assert values["raw_logits"] != values["raw_logprobs"] |
| 672 | assert values["processed_logits"] == values["raw_logits"] |
| 673 | assert values["processed_logprobs"] == values["raw_logprobs"] |
| 674 | |
| 675 | |
| 676 | class TestCorrectDecodedToken: |
| 677 | """Unit tests for _correct_decoded_token method in LogprobsProcessor. |
| 678 | |
| 679 | This method handles UTF-8 decoding issues where incomplete byte sequences |
| 680 | result in the Unicode replacement character "�" (U+FFFD). This commonly |
| 681 | happens with byte-fallback tokenization when multi-byte UTF-8 characters |
| 682 | are split across tokens. |
| 683 | |
| 684 | The method signature is _correct_decoded_token(token_id, context_token_ids) |
| 685 | where token_id is the single token to correct and context_token_ids are |
| 686 | the preceding sampled tokens in sequential order. |
| 687 | """ |
| 688 | |
| 689 | @pytest.fixture |
| 690 | def mock_tokenizer(self): |
| 691 | """Create a mock tokenizer for testing.""" |
| 692 | from unittest.mock import Mock |
| 693 | |
| 694 | tokenizer = Mock() |
| 695 | return tokenizer |
| 696 | |
| 697 | @pytest.fixture |
| 698 | def processor(self, mock_tokenizer): |
| 699 | """Create a LogprobsProcessor.""" |
| 700 | from vllm.v1.engine.logprobs import LogprobsProcessor |
| 701 | |
| 702 | processor = LogprobsProcessor( |
| 703 | tokenizer=mock_tokenizer, |
| 704 | logprobs=[], |
| 705 | prompt_logprobs=None, |
| 706 | cumulative_logprob=0.0, |
| 707 | num_logprobs=1, |
| 708 | num_prompt_logprobs=None, |
| 709 | ) |
| 710 | return processor |
| 711 | |
| 712 | def test_correction_with_context(self, processor): |
| 713 | """Test correction using context from preceding sampled tokens. |
| 714 | |
| 715 | Scenario: A byte-fallback token that completes a multi-byte |
| 716 | UTF-8 sequence when decoded with context. |
| 717 | """ |
| 718 | |
| 719 | # Context is [101] (a preceding sampled token) |
| 720 | # Token 102 individually decodes to "�" |
| 721 | # decode([101, 102]) returns "valid" (complete sequence) |
| 722 | def mock_decode(ids): |
| 723 | if ids == [101, 102]: |
| 724 | return "hello valid" |
| 725 | if ids == [101]: |
| 726 | return "hello " |
| 727 | return "�" |
| 728 | |
| 729 | processor.tokenizer.decode.side_effect = mock_decode |
| 730 | |
| 731 | result = processor._correct_decoded_token(102, [101]) |
| 732 | assert result == "valid" |
| 733 | |
| 734 | def test_correction_with_context_from_logprobs(self, processor): |
| 735 | """Test correction using context from previous logprob entries. |
| 736 | |
| 737 | Scenario: Token decoded with context from previously sampled |
| 738 | tokens completes a UTF-8 sequence. |
| 739 | """ |
| 740 | |
| 741 | # Token 123 was previously sampled (in context) |
| 742 | def mock_decode(ids): |
| 743 | if ids == [123, 100]: |
| 744 | return 'hello "polarized"' |
| 745 | if ids == [123]: |
| 746 | return "hello " |
| 747 | return "�" |
| 748 | |
| 749 | processor.tokenizer.decode.side_effect = mock_decode |
| 750 | |
| 751 | result = processor._correct_decoded_token(100, [123]) |
| 752 | assert result == '"polarized"' |
| 753 | |
| 754 | def test_correction_no_context(self, processor): |
| 755 | """Test correction with no context available. |
| 756 | |
| 757 | Should return empty string as fallback. |
| 758 | """ |
| 759 | processor.tokenizer.decode.return_value = "�" |
| 760 | |
| 761 | result = processor._correct_decoded_token(100, []) |
| 762 | assert result == "" |
| 763 | |
| 764 | def test_correction_with_context_succeeds(self, processor): |
| 765 | """Test correction with context from previously sampled tokens.""" |
| 766 | |
| 767 | def mock_decode(ids): |
| 768 | if ids == [123, 200]: |
| 769 | return "hello corrected" |
| 770 | if ids == [123]: |
| 771 | return "hello " |
| 772 | return "�" |
| 773 | |
| 774 | processor.tokenizer.decode.side_effect = mock_decode |
| 775 | |
| 776 | result = processor._correct_decoded_token(200, [123]) |
| 777 | assert result == "corrected" |
| 778 | |
| 779 | def test_fallback_when_all_attempts_fail(self, processor): |
| 780 | """Test fallback to empty string when no correction works.""" |
| 781 | processor.tokenizer.decode.return_value = "still�" |
| 782 | |
| 783 | result = processor._correct_decoded_token(102, [100, 101]) |
| 784 | assert result == "" |
| 785 | |
| 786 | def test_increasing_context_window(self, processor): |
| 787 | """Test that increasing context window finds the correction. |
| 788 | |
| 789 | Scenario: 3-byte UTF-8 char. With 1 context token, still |
| 790 | incomplete. With 2 context tokens, completes the sequence. |
| 791 | """ |
| 792 | |
| 793 | def mock_decode(ids): |
| 794 | # 1 context token: still incomplete |
| 795 | if ids == [81, 82]: |
| 796 | return "�" |
| 797 | # 2 context tokens: complete |
| 798 | if ids == [80, 81, 82]: |
| 799 | return "\u201c" |
| 800 | # Context-only decodes |
| 801 | if ids == [81]: |
| 802 | return "�" |
| 803 | if ids == [80, 81]: |
| 804 | return "�" |
| 805 | return "�" |
| 806 | |
| 807 | processor.tokenizer.decode.side_effect = mock_decode |
| 808 | |
| 809 | # Context has 2 preceding tokens [80, 81] |
| 810 | result = processor._correct_decoded_token(82, [80, 81]) |
| 811 | assert result == "\u201c" |
| 812 | |
| 813 | def test_multiple_consecutive_replacement_chars(self, processor): |
| 814 | """Test handling of multiple consecutive replacement characters. |
| 815 | |
| 816 | Scenario: Multi-byte sequence where intermediate bytes return |
| 817 | empty string and the final byte returns the complete character. |
| 818 | """ |
| 819 | processor.tokenizer.decode.return_value = "still�" |
| 820 | |
| 821 | # First byte with no useful context: returns "" |
| 822 | result1 = processor._correct_decoded_token(100, [50]) |
| 823 | assert result1 == "" |
| 824 | |
| 825 | # Second byte with same context: still returns "" |
| 826 | result2 = processor._correct_decoded_token(101, [50]) |
| 827 | assert result2 == "" |
| 828 | |
| 829 | def test_correction_with_multibyte_utf8(self, processor): |
| 830 | """Test correction involving multi-byte UTF-8 characters. |
| 831 | |
| 832 | Scenario: Byte-fallback tokenization splits curly quotes. |
| 833 | The last byte token should produce the complete character. |
| 834 | """ |
| 835 | |
| 836 | def mock_decode(ids): |
| 837 | # Context [123] + first byte: completes to left curly quote |
| 838 | if ids == [123, 200]: |
| 839 | return "hello \u201c" |
| 840 | if ids == [123]: |
| 841 | return "hello " |
| 842 | # Context [123] + second byte: completes to right curly quote |
| 843 | if ids == [123, 201]: |
| 844 | return "hello \u201d" |
| 845 | return "\ufffd" |
| 846 | |
| 847 | processor.tokenizer.decode.side_effect = mock_decode |
| 848 | |
| 849 | # Each top-k token is corrected independently with same context |
| 850 | result1 = processor._correct_decoded_token(200, [123]) |
| 851 | assert result1 == "\u201c" |
| 852 | |
| 853 | result2 = processor._correct_decoded_token(201, [123]) |
| 854 | assert result2 == "\u201d" |
| 855 | |
| 856 | def test_topk_tokens_corrected_independently(self, processor): |
| 857 | """Test that top-k alternatives at the same position are each |
| 858 | corrected independently using only sequential context, not |
| 859 | each other. |
| 860 | |
| 861 | This is the core fix for issue #27300: when logprobs > 0, |
| 862 | alternative tokens must not be combined with each other. |
| 863 | """ |
| 864 | # Context: previously sampled token 50 |
| 865 | context = [50] |
| 866 | |
| 867 | def mock_decode(ids): |
| 868 | # Token 100 (sampled) with context |
| 869 | if ids == [50, 100]: |
| 870 | return "prefix \u201c" |
| 871 | # Token 200 (top-k alternative) with context |
| 872 | if ids == [50, 200]: |
| 873 | return "prefix \u2014" |
| 874 | # Context alone |
| 875 | if ids == [50]: |
| 876 | return "prefix " |
| 877 | return "\ufffd" |
| 878 | |
| 879 | processor.tokenizer.decode.side_effect = mock_decode |
| 880 | |
| 881 | # Both tokens at the same position use the SAME context [50] |
| 882 | result_sampled = processor._correct_decoded_token(100, context) |
| 883 | assert result_sampled == "\u201c" |
| 884 | |
| 885 | result_alt = processor._correct_decoded_token(200, context) |
| 886 | assert result_alt == "\u2014" |
| 887 | |
| 888 | def test_real_world_opt125m_scenario(self, mock_tokenizer): |
| 889 | """Test the real-world scenario from the bug report. |
| 890 | |
| 891 | Simulates the OPT-125m sequence where curly quotes are split |
| 892 | into byte-fallback tokens. Each token is corrected using only |
| 893 | the preceding sampled tokens as context. |
| 894 | """ |
| 895 | from vllm.v1.engine.logprobs import LogprobsProcessor |
| 896 | |
| 897 | processor = LogprobsProcessor( |
| 898 | tokenizer=mock_tokenizer, |
| 899 | logprobs=[], |
| 900 | prompt_logprobs=None, |
| 901 | cumulative_logprob=0.0, |
| 902 | num_logprobs=1, |
| 903 | num_prompt_logprobs=None, |
| 904 | ) |
| 905 | |
| 906 | # Simulating: byte tokens 3, 4 form left curly quote "\u201c" |
| 907 | # byte tokens 8, 9 form right curly quote "\u201d" |
| 908 | def mock_decode(ids): |
| 909 | # Context decodes |
| 910 | if ids == [2]: |
| 911 | return " term" |
| 912 | if ids == [1, 2]: |
| 913 | return " the term" |
| 914 | if ids == [3]: |
| 915 | return "\ufffd" |
| 916 | if ids == [2, 3]: |
| 917 | return " term\ufffd" |
| 918 | if ids == [1, 2, 3]: |
| 919 | return " the term\ufffd" |
| 920 | # Token 4 with context [2, 3] -> completes left curly quote |
| 921 | if ids == [3, 4]: |
| 922 | return "\u201c" |
| 923 | if ids == [2, 3, 4]: |
| 924 | return " term\u201c" |
| 925 | # Context for right curly quote |
| 926 | if ids == [7]: |
| 927 | return "ized" |
| 928 | if ids == [7, 8]: |
| 929 | return "ized\ufffd" |
| 930 | if ids == [8, 9]: |
| 931 | return "\u201d" |
| 932 | if ids == [7, 8, 9]: |
| 933 | return "ized\u201d" |
| 934 | return "normal_text" |
| 935 | |
| 936 | mock_tokenizer.decode.side_effect = mock_decode |
| 937 | |
| 938 | # First byte (token 3) of left curly quote with no context |
| 939 | result = processor._correct_decoded_token(3, []) |
| 940 | assert result == "" |
| 941 | |
| 942 | # First byte (token 3) with context [2] -> still incomplete |
| 943 | result = processor._correct_decoded_token(3, [2]) |
| 944 | assert result == "" |
| 945 | |
| 946 | # Second byte (token 4) of left curly quote with context [2, 3] |
| 947 | # Token 3 is byte-fallback, so clean context is [2] only. |
| 948 | # decode([2, 3, 4]) = " term\u201c", decode([2]) = " term" |
| 949 | # result = "\u201c" |
| 950 | result = processor._correct_decoded_token(4, [2, 3]) |
| 951 | assert result == "\u201c" |
| 952 | |
| 953 | # Second byte (token 9) of right curly quote with context [7, 8] |
| 954 | result = processor._correct_decoded_token(9, [7, 8]) |
| 955 | assert result == "\u201d" |
| 956 | |
| 957 | def test_byte_fallback_context_preserves_space(self, mock_tokenizer): |
| 958 | """Test that text from byte-fallback context tokens is preserved. |
| 959 | |
| 960 | In OPT-125m, token 44 = space + 2 bytes of curly quote. |
| 961 | When token 44 returns "" (incomplete), the space it carried |
| 962 | must be attributed to the completing token (48). |
| 963 | """ |
| 964 | from vllm.v1.engine.logprobs import LogprobsProcessor |
| 965 | |
| 966 | processor = LogprobsProcessor( |
| 967 | tokenizer=mock_tokenizer, |
| 968 | logprobs=[], |
| 969 | prompt_logprobs=None, |
| 970 | cumulative_logprob=0.0, |
| 971 | num_logprobs=1, |
| 972 | num_prompt_logprobs=None, |
| 973 | ) |
| 974 | |
| 975 | def mock_decode(ids): |
| 976 | # Token 44 = space + 2 bytes (like OPT-125m's \u0120\u00e2\u0080) |
| 977 | if ids == [44]: |
| 978 | return " \ufffd" |
| 979 | if ids == [48]: |
| 980 | return "\ufffd" |
| 981 | # Together they form: space + left curly quote |
| 982 | if ids == [44, 48]: |
| 983 | return " \u201c" |
| 984 | # With preceding clean context |
| 985 | if ids == [1385]: |
| 986 | return " term" |
| 987 | if ids == [1385, 44]: |
| 988 | return " term \ufffd" |
| 989 | if ids == [1385, 44, 48]: |
| 990 | return " term \u201c" |
| 991 | return "\ufffd" |
| 992 | |
| 993 | mock_tokenizer.decode.side_effect = mock_decode |
| 994 | |
| 995 | # Token 44 with context [1385] -> still ends with replacement |
| 996 | result = processor._correct_decoded_token(44, [1385]) |
| 997 | assert result == "" |
| 998 | |
| 999 | # Token 48 with context [1385, 44]: |
| 1000 | # Token 44 is byte-fallback, so clean context is [1385]. |
| 1001 | # decode([1385, 44, 48]) = " term \u201c" |
| 1002 | # decode([1385]) = " term" |
| 1003 | # result = " \u201c" (space preserved from token 44!) |
| 1004 | result = processor._correct_decoded_token(48, [1385, 44]) |
| 1005 | assert result == " \u201c" |
| 1006 | |
| 1007 | |
| 1008 | def test_verify_tokens_integration(): |
| 1009 | """Integration test for _verify_tokens with real model. |
| 1010 | |
| 1011 | This test validates that _verify_tokens correctly identifies and |
| 1012 | corrects tokens ending with the replacement character "�". |
| 1013 | Uses facebook/opt-125m which is known to produce these issues. |
| 1014 | """ |
| 1015 | with VllmRunner( |
| 1016 | "facebook/opt-125m", |
| 1017 | max_logprobs=0, |
| 1018 | enable_prefix_caching=False, |
| 1019 | gpu_memory_utilization=0.15, |
| 1020 | max_model_len=256, |
| 1021 | ) as runner: |
| 1022 | # Use a prompt that triggers multi-byte UTF-8 issues |
| 1023 | # Based on user's example: "In this example," |
| 1024 | test_prompts = ["In this example,"] |
| 1025 | |
| 1026 | sampling_params = SamplingParams( |
| 1027 | max_tokens=16, |
| 1028 | temperature=0, |
| 1029 | logprobs=0, |
| 1030 | ) |
| 1031 | |
| 1032 | results = runner.llm.generate(test_prompts, sampling_params=sampling_params) |
| 1033 | |
| 1034 | # Verify that decoded tokens don't contain replacement characters |
| 1035 | for result in results: |
| 1036 | assert result.outputs[0].logprobs is not None |
| 1037 | for logprob_dict in result.outputs[0].logprobs: |
| 1038 | for token_id, logprob_info in logprob_dict.items(): |
| 1039 | decoded_token = logprob_info.decoded_token |
| 1040 | # Decoded tokens should not end with replacement character |
| 1041 | # They should either be corrected or empty string |
| 1042 | assert not decoded_token.endswith("�"), ( |
| 1043 | f"Token {token_id} decoded to '{decoded_token}' which " |
| 1044 | f"ends with replacement character" |
| 1045 | ) |
| 1046 | # Decoded tokens should not contain lone replacement characters |
| 1047 | assert decoded_token != "�", ( |
| 1048 | f"Token {token_id} is a lone replacement character" |
| 1049 | ) |
| 1050 | |
| 1051 | |
| 1052 | def test_utf8_edge_cases_with_real_model(): |
| 1053 | """Test various UTF-8 edge cases with a real model. |
| 1054 | |
| 1055 | Tests prompts that are likely to trigger byte-fallback tokenization |
| 1056 | and multi-byte UTF-8 splitting. |
| 1057 | """ |
| 1058 | with VllmRunner( |
| 1059 | "facebook/opt-125m", |
| 1060 | max_logprobs=1, |
| 1061 | enable_prefix_caching=False, |
| 1062 | gpu_memory_utilization=0.15, |
| 1063 | max_model_len=256, |
| 1064 | ) as runner: |
| 1065 | # Prompts with various multi-byte UTF-8 characters |
| 1066 | test_prompts = [ |
| 1067 | 'Smart quotes: "Hello"', # Curly quotes |
| 1068 | "Em dash — test", # Em dash |
| 1069 | "Ellipsis… continues", # Ellipsis |
| 1070 | "Chinese: 你好", # Chinese characters |
| 1071 | "Emoji: 😀 🎉", # Emojis |
| 1072 | 'Mixed: "quoted" — with symbols', # Mixed |
| 1073 | ] |
| 1074 | |
| 1075 | sampling_params = SamplingParams( |
| 1076 | max_tokens=10, |
| 1077 | temperature=0, |
| 1078 | logprobs=1, |
| 1079 | ) |
| 1080 | |
| 1081 | results = runner.llm.generate(test_prompts, sampling_params=sampling_params) |
| 1082 | |
| 1083 | for i, result in enumerate(results): |
| 1084 | prompt = test_prompts[i] |
| 1085 | assert result.outputs[0].logprobs is not None |
| 1086 | |
| 1087 | # Check that no decoded tokens end with replacement character |
| 1088 | for logprob_dict in result.outputs[0].logprobs: |
| 1089 | for token_id, logprob_info in logprob_dict.items(): |
| 1090 | decoded_token = logprob_info.decoded_token |
| 1091 | assert not decoded_token.endswith("�"), ( |
| 1092 | f"Prompt: '{prompt}'\n" |
| 1093 | f"Token {token_id} decoded to '{decoded_token}' which " |
| 1094 | f"ends with replacement character" |
| 1095 | ) |
| 1096 | |
| 1097 | |
| 1098 | def test_correct_decoded_token_preserves_valid_tokens(): |
| 1099 | """Test that valid tokens (not ending with �) are not modified. |
| 1100 | |
| 1101 | The _correct_decoded_token method should only be called for tokens |
| 1102 | ending with "�", but this test verifies the broader _verify_tokens |
| 1103 | logic doesn't affect valid tokens. |
| 1104 | """ |
| 1105 | with VllmRunner( |
| 1106 | "facebook/opt-125m", |
| 1107 | max_logprobs=2, |
| 1108 | enable_prefix_caching=False, |
| 1109 | gpu_memory_utilization=0.15, |
| 1110 | max_model_len=256, |
| 1111 | ) as runner: |
| 1112 | # Simple prompt with standard ASCII characters |
| 1113 | test_prompts = ["Hello world, this is a test."] |
| 1114 | |
| 1115 | sampling_params = SamplingParams( |
| 1116 | max_tokens=10, |
| 1117 | temperature=0, |
| 1118 | logprobs=2, |
| 1119 | ) |
| 1120 | |
| 1121 | results = runner.llm.generate(test_prompts, sampling_params=sampling_params) |
| 1122 | |
| 1123 | for result in results: |
| 1124 | assert result.outputs[0].logprobs is not None |
| 1125 | |
| 1126 | # All decoded tokens should be valid strings |
| 1127 | for logprob_dict in result.outputs[0].logprobs: |
| 1128 | for token_id, logprob_info in logprob_dict.items(): |
| 1129 | decoded_token = logprob_info.decoded_token |
| 1130 | # Valid tokens should be non-empty strings (or empty if corrected) |
| 1131 | assert isinstance(decoded_token, str) |
| 1132 | # Should not contain replacement character |
| 1133 | assert "�" not in decoded_token |
| 1134 | |
| 1135 | |
| 1136 | @pytest.mark.parametrize("logprobs_mode", get_args(LogprobsMode)) |
| 1137 | @pytest.mark.parametrize( |
| 1138 | "model_setup", |
| 1139 | [ |
| 1140 | pytest.param( |
| 1141 | ( |
| 1142 | "eagle", |
| 1143 | "meta-llama/Llama-3.2-1B-Instruct", |
| 1144 | { |
| 1145 | "method": "eagle", |
| 1146 | "model": "nm-testing/Llama3_2_1B_speculator.eagle3", |
| 1147 | "num_speculative_tokens": 3, |
| 1148 | }, |
| 1149 | 0, |
| 1150 | ), |
| 1151 | marks=large_gpu_mark(min_gb=32), |
| 1152 | id="eagle0", |
| 1153 | ), |
| 1154 | pytest.param( |
| 1155 | ( |
| 1156 | "eagle", |
| 1157 | "meta-llama/Llama-3.2-1B-Instruct", |
| 1158 | { |
| 1159 | "method": "eagle", |
| 1160 | "model": "nm-testing/Llama3_2_1B_speculator.eagle3", |
| 1161 | "num_speculative_tokens": 3, |
| 1162 | }, |
| 1163 | 3, |
| 1164 | ), |
| 1165 | marks=large_gpu_mark(min_gb=32), |
| 1166 | id="eagle3", |
| 1167 | ), |
| 1168 | pytest.param( |
| 1169 | ( |
| 1170 | "ngram", |
| 1171 | "meta-llama/Llama-3.2-1B-Instruct", |
| 1172 | { |
| 1173 | "method": "ngram", |
| 1174 | "prompt_lookup_max": 5, |
| 1175 | "prompt_lookup_min": 3, |
| 1176 | "num_speculative_tokens": 3, |
| 1177 | }, |
| 1178 | 3, |
| 1179 | ), |
| 1180 | marks=large_gpu_mark(min_gb=32), |
| 1181 | id="ngram", |
| 1182 | ), |
| 1183 | ], |
| 1184 | ) |
| 1185 | def test_spec_decode_logprobs( |
| 1186 | logprobs_mode: LogprobsMode, |
| 1187 | model_setup: tuple[str, str, dict, int], |
| 1188 | ): |
| 1189 | """Spec decode logprobs should match those of the base model. |
| 1190 | |
| 1191 | Runs the base model and spec decode model sequentially, ensuring |
| 1192 | only one LLM instance is alive at a time to avoid GPU memory |
| 1193 | contention. Both use identical chunked prefill settings and eager |
| 1194 | mode to control for infrastructure differences. |
| 1195 | |
| 1196 | Args: |
| 1197 | logprobs_mode: logprobs mode. |
| 1198 | model_setup: Tuple of (method, base model name, |
| 1199 | speculative_config dict, top_logprobs). |
| 1200 | |
| 1201 | """ |
| 1202 | from vllm import LLM |
| 1203 | |
| 1204 | method, model_name, spec_config, top_logprobs = model_setup |
| 1205 | |
| 1206 | prompt = "Hello world " * 50 |
| 1207 | sampling_params = SamplingParams( |
| 1208 | temperature=0, logprobs=top_logprobs, max_tokens=10, ignore_eos=False |
| 1209 | ) |
| 1210 | penalty_sampling_params = SamplingParams( |
| 1211 | temperature=0, |
| 1212 | logprobs=top_logprobs, |
| 1213 | max_tokens=10, |
| 1214 | ignore_eos=False, |
| 1215 | presence_penalty=-1.0, |
| 1216 | ) |
| 1217 | |
| 1218 | max_model_len = 256 |
| 1219 | llm_kwargs = dict( |
| 1220 | max_logprobs=5, |
| 1221 | max_model_len=max_model_len, |
| 1222 | seed=42, |
| 1223 | logprobs_mode=logprobs_mode, |
| 1224 | gpu_memory_utilization=0.4, |
| 1225 | # Force the same prefill chunking for both the base model and |
| 1226 | # spec decode model so the comparison isolates spec decode. |
| 1227 | enable_chunked_prefill=True, |
| 1228 | max_num_batched_tokens=32, |
| 1229 | enable_prefix_caching=False, |
| 1230 | **GPU_DETERMINISM_KWARGS, |
| 1231 | ) |
| 1232 | |
| 1233 | # Run base LLM. |
| 1234 | ref_llm = LLM( |
| 1235 | model=model_name, |
| 1236 | **llm_kwargs, |
| 1237 | ) |
| 1238 | ref_results = ref_llm.generate( |
| 1239 | [prompt, prompt], [sampling_params, penalty_sampling_params] |
| 1240 | ) |
| 1241 | # Collect logprobs outputs from reference LLM. |
| 1242 | ref_logprobs: list[Logprob] = [] |
| 1243 | for results in ref_results: |
| 1244 | for output in results.outputs: |
| 1245 | assert output.logprobs is not None |
| 1246 | for logprobs in output.logprobs: |
| 1247 | assert logprobs is not None |
| 1248 | ref_logprobs.extend(logprobs.values()) |
| 1249 | del ref_llm |
| 1250 | torch.accelerator.empty_cache() |
| 1251 | cleanup_dist_env_and_memory() |
| 1252 | |
| 1253 | # Run spec decode LLM. |
| 1254 | # Add max_model_len to spec_config if not present |
| 1255 | spec_config_with_len = {**spec_config, "max_model_len": max_model_len} |
| 1256 | spec_llm = LLM( |
| 1257 | model_name, |
| 1258 | speculative_config=spec_config_with_len, |
| 1259 | **llm_kwargs, |
| 1260 | ) |
| 1261 | spec_results = spec_llm.generate( |
| 1262 | [prompt, prompt], [sampling_params, penalty_sampling_params] |
| 1263 | ) |
| 1264 | # Collect logprobs outputs from spec decode LLM. |
| 1265 | spec_logprobs: list[Logprob] = [] |
| 1266 | for results in spec_results: |
| 1267 | for output in results.outputs: |
| 1268 | assert output.logprobs is not None |
| 1269 | for logprobs in output.logprobs: |
| 1270 | assert logprobs is not None |
| 1271 | spec_logprobs.extend(logprobs.values()) |
| 1272 | del spec_llm |
| 1273 | torch.accelerator.empty_cache() |
| 1274 | cleanup_dist_env_and_memory() |
| 1275 | |
| 1276 | # Per-token logprobs are expected to be the same. |
| 1277 | assert len(ref_logprobs) == len(spec_logprobs) |
| 1278 | for ref_logprob, spec_logprob in zip(ref_logprobs, spec_logprobs): |
| 1279 | assert math.isclose( |
| 1280 | ref_logprob.logprob, spec_logprob.logprob, rel_tol=5e-2, abs_tol=2.5e-1 |
| 1281 | ), ( |
| 1282 | f"Logprob mismatch: ref={ref_logprob.logprob} " |
| 1283 | f"spec={spec_logprob.logprob} " |
| 1284 | f"diff={abs(ref_logprob.logprob - spec_logprob.logprob)} " |
| 1285 | f"(token={ref_logprob.decoded_token!r})" |
| 1286 | ) |
| 1287 | assert ref_logprob.rank == spec_logprob.rank, ( |
| 1288 | f"Rank mismatch: ref={ref_logprob.rank} " |
| 1289 | f"spec={spec_logprob.rank} " |
| 1290 | f"(token={ref_logprob.decoded_token!r})" |
| 1291 | ) |
| 1292 | assert ref_logprob.decoded_token == spec_logprob.decoded_token |
| 1293 | |
| 1294 | |
| 1295 | def test_prompt_logprobs_with_chunking_and_preemption(): |
| 1296 | """Test that prompt logprobs are correctly returned when using |
| 1297 | both chunked prefill and preemption. |
| 1298 | |
| 1299 | This test ensures that the num_prompt_logprobs tracking persists |
| 1300 | across preemptions and prefill chunks. |
| 1301 | """ |
| 1302 | # Create prompts that will trigger chunking and preemption |
| 1303 | prompts = [ |
| 1304 | "The following numbers of the sequence " |
| 1305 | + ", ".join(str(i) for i in range(10)) |
| 1306 | + " are:", |
| 1307 | "In one word, the capital of France is ", |
| 1308 | ] + [f"Tell me about the number {i}: " for i in range(32)] |
| 1309 | |
| 1310 | sampling_params = SamplingParams( |
| 1311 | temperature=0.0, |
| 1312 | max_tokens=40, |
| 1313 | min_tokens=20, |
| 1314 | prompt_logprobs=2, # Request prompt logprobs |
| 1315 | ) |
| 1316 | |
| 1317 | with VllmRunner( |
| 1318 | "Qwen/Qwen3-0.6B", |
| 1319 | max_model_len=512, |
| 1320 | enable_chunked_prefill=True, |
| 1321 | max_num_batched_tokens=48, # Force prefill chunking |
| 1322 | num_gpu_blocks_override=33, # Force preemptions (32 usable + null block) |
| 1323 | disable_log_stats=False, |
| 1324 | gpu_memory_utilization=0.25, |
| 1325 | ) as vllm_model: |
| 1326 | metrics_before = vllm_model.llm.get_metrics() |
| 1327 | |
| 1328 | # Generate with prompt logprobs using generate_w_logprobs which |
| 1329 | # returns (output_ids, output_str, output_logprobs, prompt_logprobs) |
| 1330 | outputs = vllm_model.generate_w_logprobs( |
| 1331 | prompts, sampling_params=sampling_params, include_prompt_token_ids=True |
| 1332 | ) |
| 1333 | |
| 1334 | # Verify that all outputs have prompt logprobs |
| 1335 | for i, output in enumerate(outputs): |
| 1336 | _, _, _, prompt_token_ids, prompt_logprobs = output |
| 1337 | assert prompt_logprobs is not None and len(prompt_logprobs) > 0, ( |
| 1338 | f"Output {i} missing prompt logprobs" |
| 1339 | ) |
| 1340 | assert len(prompt_logprobs) == len(prompt_token_ids), ( |
| 1341 | "Unexpected number of prompt logprob positions" |
| 1342 | ) |
| 1343 | |
| 1344 | assert sampling_params.prompt_logprobs is not None |
| 1345 | # Each position should have the requested number of logprobs |
| 1346 | for pos, logprobs_dict in enumerate(prompt_logprobs): |
| 1347 | if logprobs_dict is not None: # First token may be None |
| 1348 | assert ( |
| 1349 | sampling_params.prompt_logprobs |
| 1350 | <= len(logprobs_dict) |
| 1351 | <= sampling_params.prompt_logprobs + 1 |
| 1352 | ), ( |
| 1353 | f"Output {i} position {pos} has {len(logprobs_dict)} " |
| 1354 | f"logprobs, expected {sampling_params.prompt_logprobs}" |
| 1355 | ) |
| 1356 | |
| 1357 | # Check that we actually had preemptions |
| 1358 | metrics_after = vllm_model.llm.get_metrics() |
| 1359 | preemptions_before = next( |
| 1360 | (m.value for m in metrics_before if m.name == "vllm:num_preemptions"), 0 |
| 1361 | ) |
| 1362 | preemptions_after = next( |
| 1363 | (m.value for m in metrics_after if m.name == "vllm:num_preemptions"), 0 |
| 1364 | ) |
| 1365 | preemptions = preemptions_after - preemptions_before |
| 1366 | assert preemptions > 0, "Test did not trigger any preemptions" |
| 1367 | |
| 1368 | print(f"Test passed with {preemptions} preemptions") |
| 1369 | |
| 1370 | |
| 1371 | def test_prompt_logprob_token_ids_with_chunking_and_preemption(monkeypatch): |
| 1372 | """Fixed-ID scores stay row-aligned across chunked prefill and preemption.""" |
| 1373 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 1374 | |
| 1375 | prompts = [ |
| 1376 | "The following numbers of the sequence " |
| 1377 | + ", ".join(str(i) for i in range(10)) |
| 1378 | + " are:", |
| 1379 | "In one word, the capital of France is ", |
| 1380 | ] + [f"Tell me about the number {i}: " for i in range(32)] |
| 1381 | |
| 1382 | start = 2 |
| 1383 | candidate_ids = [10, 100, 1000, 10000] |
| 1384 | sampling_params = SamplingParams( |
| 1385 | temperature=0.0, |
| 1386 | max_tokens=40, |
| 1387 | min_tokens=20, |
| 1388 | prompt_logprob_token_ids=candidate_ids, |
| 1389 | prompt_logprob_start=start, |
| 1390 | ) |
| 1391 | |
| 1392 | with VllmRunner( |
| 1393 | "Qwen/Qwen3-0.6B", |
| 1394 | max_model_len=512, |
| 1395 | enable_chunked_prefill=True, |
| 1396 | max_num_batched_tokens=48, # Force prefill chunking |
| 1397 | num_gpu_blocks_override=33, # Force preemptions (32 usable + null block) |
| 1398 | disable_log_stats=False, |
| 1399 | gpu_memory_utilization=0.25, |
| 1400 | ) as vllm_model: |
| 1401 | metrics_before = vllm_model.llm.get_metrics() |
| 1402 | outputs = vllm_model.llm.generate(prompts, sampling_params) |
| 1403 | |
| 1404 | for i, output in enumerate(outputs): |
| 1405 | scores = output.prompt_token_id_logprobs |
| 1406 | assert scores is not None, f"Output {i} missing fixed-ID scores" |
| 1407 | expected_shape = ( |
| 1408 | len(output.prompt_token_ids) - 1 - start, |
| 1409 | len(candidate_ids), |
| 1410 | ) |
| 1411 | assert scores.shape == expected_shape, ( |
| 1412 | f"Output {i} scored {scores.shape}, expected {expected_shape}" |
| 1413 | ) |
| 1414 | assert math.isfinite(float(scores.min())) |
| 1415 | assert float(scores.max()) <= 1e-3, "logprobs must be <= 0" |
| 1416 | |
| 1417 | metrics_after = vllm_model.llm.get_metrics() |
| 1418 | preemptions_before = next( |
| 1419 | (m.value for m in metrics_before if m.name == "vllm:num_preemptions"), 0 |
| 1420 | ) |
| 1421 | preemptions_after = next( |
| 1422 | (m.value for m in metrics_after if m.name == "vllm:num_preemptions"), 0 |
| 1423 | ) |
| 1424 | assert preemptions_after - preemptions_before > 0, ( |
| 1425 | "Test did not trigger any preemptions" |
| 1426 | ) |
| 1427 | |
| 1428 | # Row alignment is numerical: chunked and preempted scores must match an |
| 1429 | # unchunked, unpreempted run of the same requests. Batch composition moves |
| 1430 | # bf16 tail logprobs by a few percent; a misaligned row differs by nats. |
| 1431 | with VllmRunner( |
| 1432 | "Qwen/Qwen3-0.6B", |
| 1433 | max_model_len=512, |
| 1434 | max_logprobs=64, |
| 1435 | gpu_memory_utilization=0.25, |
| 1436 | ) as reference: |
| 1437 | reference_outputs = reference.llm.generate(prompts, sampling_params) |
| 1438 | own_ids = [list(dict.fromkeys(o.prompt_token_ids)) for o in reference_outputs] |
| 1439 | self_scored = reference.llm.generate( |
| 1440 | prompts, |
| 1441 | [ |
| 1442 | SamplingParams( |
| 1443 | max_tokens=1, prompt_logprobs=0, prompt_logprob_token_ids=ids |
| 1444 | ) |
| 1445 | for ids in own_ids |
| 1446 | ], |
| 1447 | ) |
| 1448 | for output, ref in zip(outputs, reference_outputs): |
| 1449 | np.testing.assert_allclose( |
| 1450 | output.prompt_token_id_logprobs, ref.prompt_token_id_logprobs, rtol=0.1 |
| 1451 | ) |
| 1452 | for output, ids in zip(self_scored, own_ids): |
| 1453 | scores = output.prompt_token_id_logprobs |
| 1454 | for row, target in enumerate(output.prompt_token_ids[1:]): |
| 1455 | expected = output.prompt_logprobs[row + 1][target].logprob |
| 1456 | assert scores[row, ids.index(target)] == pytest.approx(expected, abs=1e-3) |
| 1457 | |
| 1458 | |
| 1459 | def test_prompt_logprob_token_ids_drop_partially_scored_prefills(monkeypatch): |
| 1460 | """A prefill that starts past the first scored row returns no scores. |
| 1461 | |
| 1462 | Such a prefill (here a prefix-cache hit the caller opted back into) never |
| 1463 | computes the leading rows, so emitting the buffer would return values the |
| 1464 | model never produced for them. |
| 1465 | """ |
| 1466 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1") |
| 1467 | |
| 1468 | prompt = "The capital of France is Paris. " * 20 |
| 1469 | sampling_params = SamplingParams( |
| 1470 | temperature=0.0, |
| 1471 | max_tokens=1, |
| 1472 | prompt_logprob_token_ids=[10, 100, 1000], |
| 1473 | skip_reading_prefix_cache=False, |
| 1474 | ) |
| 1475 | |
| 1476 | with VllmRunner( |
| 1477 | "Qwen/Qwen3-0.6B", |
| 1478 | max_model_len=512, |
| 1479 | enable_prefix_caching=True, |
| 1480 | gpu_memory_utilization=0.25, |
| 1481 | ) as vllm_model: |
| 1482 | first = vllm_model.llm.generate([prompt], sampling_params)[0] |
| 1483 | assert first.prompt_token_id_logprobs is not None |
| 1484 | assert math.isfinite(float(first.prompt_token_id_logprobs.min())) |
| 1485 | |
| 1486 | # The re-send is served from the prefix cache, so its prefill starts |
| 1487 | # past row 0 and the leading rows are never computed. |
| 1488 | second = vllm_model.llm.generate([prompt], sampling_params)[0] |
| 1489 | assert second.prompt_token_id_logprobs is None |
| 1490 | |
| 1491 | |
| 1492 | @large_gpu_mark(min_gb=24) |
| 1493 | def test_token_logprobs_large_batch_int64_row_offset(): |
| 1494 | """Regression: logprob kernel row offset (row * vocab_size) must use int64. |
| 1495 | |
| 1496 | The rejection-sampler logprobs path runs the logprob kernels over the |
| 1497 | spec-expanded logits batch, so batch_size * vocab_size can exceed 2**31 |
| 1498 | (e.g. DFlash drafts K tokens per request). With int32 offset arithmetic the |
| 1499 | per-row pointer wraps to a negative address and the kernel hits a CUDA |
| 1500 | illegal memory access. Run over a batch where batch_size * vocab_size > 2**31 |
| 1501 | and check the highest-offset row matches a reference log-softmax. |
| 1502 | """ |
| 1503 | if not current_platform.is_cuda(): |
| 1504 | pytest.skip("int32 row-offset overflow is a CUDA kernel issue") |
| 1505 | from vllm.v1.worker.gpu.sample.logprob import compute_token_logprobs |
| 1506 | |
| 1507 | device = torch.device("cuda") |
| 1508 | vocab_size = 131072 |
| 1509 | batch_size = 2**31 // vocab_size + 64 # batch_size * vocab_size > 2**31 |
| 1510 | # logits (the large input) plus small logprob/rank outputs; ~1 GB headroom. |
| 1511 | required_bytes = batch_size * vocab_size * 4 + (1 << 30) |
| 1512 | if torch.accelerator.get_memory_info()[0] < required_bytes: |
| 1513 | pytest.skip(f"needs ~{required_bytes / 1e9:.0f} GB of free GPU memory") |
| 1514 | |
| 1515 | logits = torch.randn(batch_size, vocab_size, device=device, dtype=torch.float32) |
| 1516 | token_ids = torch.full((batch_size, 1), 7, device=device, dtype=torch.int64) |
| 1517 | logprobs = compute_token_logprobs(logits, token_ids) |
| 1518 | torch.accelerator.synchronize() # surface any async illegal memory access |
| 1519 | last = batch_size - 1 |
| 1520 | ref = torch.log_softmax(logits[last].float(), dim=-1)[7] |
| 1521 | assert torch.allclose(logprobs[last, 0], ref, atol=1e-2), ( |
| 1522 | f"logprob {logprobs[last, 0].item()} != ref {ref.item()}" |
| 1523 | ) |