| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Numerical checks against an independent, teacher-forced target model.""" |
| 4 | |
| 5 | import json |
| 6 | import math |
| 7 | from collections import Counter |
| 8 | from dataclasses import dataclass |
| 9 | from functools import lru_cache |
| 10 | |
| 11 | import torch |
| 12 | import xgrammar as xgr |
| 13 | from tokenizers.decoders import DecodeStream |
| 14 | |
| 15 | from tests.conftest import HfRunner |
| 16 | from vllm import SamplingParams |
| 17 | from vllm.logprobs import Logprob |
| 18 | from vllm.outputs import RequestOutput |
| 19 | from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens |
| 20 | from vllm.v1.engine.logprobs import LogprobsProcessor |
| 21 | |
| 22 | |
| 23 | @dataclass(frozen=True) |
| 24 | class AccuracyTolerance: |
| 25 | logprob_atol: float |
| 26 | greedy_atol: float |
| 27 | logprob_mean_atol: float = 0.1 |
| 28 | greedy_total_atol: float = 0.5 |
| 29 | |
| 30 | def logprob_error(self, reference: torch.Tensor) -> torch.Tensor: |
| 31 | return torch.full_like(reference, self.logprob_atol) |
| 32 | |
| 33 | |
| 34 | def check_first_divergence( |
| 35 | baseline: RequestOutput, |
| 36 | actual: RequestOutput, |
| 37 | params: SamplingParams, |
| 38 | tolerance: AccuracyTolerance, |
| 39 | context: str, |
| 40 | ) -> None: |
| 41 | """Require the two engines' scores to justify their first different choice.""" |
| 42 | assert baseline.finished and actual.finished, context |
| 43 | assert len(baseline.outputs) == len(actual.outputs) == 1, context |
| 44 | assert baseline.prompt_token_ids is not None, context |
| 45 | assert baseline.prompt_token_ids == actual.prompt_token_ids, context |
| 46 | base, other = baseline.outputs[0], actual.outputs[0] |
| 47 | position = next( |
| 48 | (i for i, (a, b) in enumerate(zip(base.token_ids, other.token_ids)) if a != b), |
| 49 | None, |
| 50 | ) |
| 51 | if position is None: |
| 52 | assert len(base.token_ids) == len(other.token_ids), ( |
| 53 | f"{context}: one completion is a strict prefix of the other" |
| 54 | ) |
| 55 | assert (base.finish_reason, base.stop_reason) == ( |
| 56 | other.finish_reason, |
| 57 | other.stop_reason, |
| 58 | ), f"{context}: identical tokens have different stopping reasons" |
| 59 | return |
| 60 | |
| 61 | assert params.logprobs is not None and params.logprobs > 0, context |
| 62 | assert params.temperature == 0.0, context |
| 63 | assert params.presence_penalty == 0.0 and params.repetition_penalty == 1.0 |
| 64 | assert not params.logit_bias, context |
| 65 | a, b = base.token_ids[position], other.token_ids[position] |
| 66 | counts = Counter(base.token_ids[:position]) |
| 67 | context = ( |
| 68 | f"{context}, first divergence position={position}, " |
| 69 | f"baseline token={a}, actual token={b}" |
| 70 | ) |
| 71 | for side, completion, selected, alternative in ( |
| 72 | ("baseline", base, a, b), |
| 73 | ("actual", other, b, a), |
| 74 | ): |
| 75 | assert completion.logprobs is not None, context |
| 76 | assert len(completion.logprobs) == len(completion.token_ids), context |
| 77 | scores = completion.logprobs[position] |
| 78 | assert selected in scores, f"{context}, {side}: missing selected token score" |
| 79 | chosen = scores[selected].logprob |
| 80 | assert math.isfinite(chosen), f"{context}, {side}: non-finite selected score" |
| 81 | known = alternative in scores |
| 82 | if known: |
| 83 | competing = scores[alternative].logprob |
| 84 | assert math.isfinite(competing), ( |
| 85 | f"{context}, {side}: divergent token has a non-finite score" |
| 86 | ) |
| 87 | else: |
| 88 | # The sampled token may be returned in addition to the requested |
| 89 | # top-k. An omitted token's raw score is bounded by the kth score. |
| 90 | top_scores = [ |
| 91 | score.logprob |
| 92 | for token, score in scores.items() |
| 93 | if len(scores) == params.logprobs or token != selected |
| 94 | ] |
| 95 | assert len(top_scores) == params.logprobs, context |
| 96 | assert not any(math.isnan(score) for score in top_scores), context |
| 97 | competing = min(top_scores) |
| 98 | |
| 99 | # Both candidates were emitted with the same history and constraints; |
| 100 | # frequency penalties are the only supported transform changing their |
| 101 | # finite score difference. Prompt tokens do not contribute to counts. |
| 102 | chosen -= params.frequency_penalty * counts[selected] |
| 103 | competing -= params.frequency_penalty * counts[alternative] |
| 104 | gap = chosen - competing |
| 105 | detail = ( |
| 106 | f"{context}, {side}: selected={selected}, alternative={alternative}, " |
| 107 | f"{'gap' if known else 'gap lower bound'}={gap}" |
| 108 | ) |
| 109 | if known: |
| 110 | # Log-softmax and reconstructing penalties may round in FP32, but |
| 111 | # a near tie cannot excuse selecting the engine's own worse token. |
| 112 | assert gap >= -1e-5, f"{detail}: selected token is below its alternative" |
| 113 | assert gap <= tolerance.greedy_atol, ( |
| 114 | f"{detail}: divergence exceeds near-tie bound={tolerance.greedy_atol}" |
| 115 | ) |
| 116 | |
| 117 | |
| 118 | def check_logprobs( |
| 119 | scores: dict[int, Logprob], |
| 120 | reference: torch.Tensor, |
| 121 | token: int, |
| 122 | requested: int, |
| 123 | tolerance: AccuracyTolerance, |
| 124 | context: str, |
| 125 | ) -> torch.Tensor: |
| 126 | """Check every score, reported rank, and the complete requested top-k.""" |
| 127 | assert all(0 <= token_id < reference.numel() for token_id in scores), context |
| 128 | assert token in scores, f"{context}: missing score for token {token}" |
| 129 | assert max(1, requested) <= len(scores) <= requested + 1, context |
| 130 | assert not torch.isnan(reference).any(), context |
| 131 | assert all(s.rank is not None and s.rank >= 1 for s in scores.values()), context |
| 132 | ids = torch.tensor(list(scores), device=reference.device) |
| 133 | actual = torch.tensor( |
| 134 | [score.logprob for score in scores.values()], device=reference.device |
| 135 | ) |
| 136 | expected = reference[ids] |
| 137 | finite = torch.isfinite(expected) |
| 138 | assert torch.equal(torch.isfinite(actual), finite), ( |
| 139 | f"{context}: finite/masked logprobs differ: {scores}" |
| 140 | ) |
| 141 | assert torch.equal(actual[~finite], expected[~finite]), context |
| 142 | assert (actual[finite] <= 0).all(), f"{context}: positive log probability" |
| 143 | errors = (actual[finite] - expected[finite]).abs() |
| 144 | bounds = tolerance.logprob_error(expected[finite]) |
| 145 | assert (errors <= bounds).all(), ( |
| 146 | f"{context}: token IDs={ids[finite].tolist()}, " |
| 147 | f"actual={actual[finite].tolist()}, reference={expected[finite].tolist()}, " |
| 148 | f"error={errors.tolist()}, bound={bounds.tolist()}" |
| 149 | ) |
| 150 | |
| 151 | for token_id, score in scores.items(): |
| 152 | value = reference[token_id] |
| 153 | if not torch.isfinite(value): |
| 154 | continue |
| 155 | delta = tolerance.logprob_error(value) |
| 156 | lower = 1 + int((reference > value + 2 * delta).sum()) |
| 157 | upper = max(1, int((reference >= value - 2 * delta).sum())) |
| 158 | assert score.rank is not None and lower <= score.rank <= upper, ( |
| 159 | f"{context}: token={token_id}, rank={score.rank}, " |
| 160 | f"reference rank interval=[{lower}, {upper}]" |
| 161 | ) |
| 162 | |
| 163 | if requested: |
| 164 | # The sampled token can be outside top-k even when it ties the kth |
| 165 | # score. Its strict-greater rank can therefore duplicate a top-k rank. |
| 166 | top_ids = ( |
| 167 | list(scores) |
| 168 | if len(scores) == requested |
| 169 | else [t for t in scores if t != token] |
| 170 | ) |
| 171 | assert len(top_ids) == requested, f"{context}: incomplete top-k" |
| 172 | assert {scores[t].rank for t in top_ids} == set(range(1, requested + 1)), ( |
| 173 | context |
| 174 | ) |
| 175 | ordered = sorted(top_ids, key=lambda t: scores[t].rank) |
| 176 | assert all( |
| 177 | scores[a].logprob >= scores[b].logprob for a, b in zip(ordered, ordered[1:]) |
| 178 | ), f"{context}: top-k scores disagree with their reported ranks" |
| 179 | boundary = reference.topk(requested).values[-1] |
| 180 | if torch.isfinite(boundary): |
| 181 | delta = tolerance.logprob_error(boundary) |
| 182 | required = (reference > boundary + 2 * delta).nonzero().flatten() |
| 183 | assert set(required.tolist()) <= scores.keys(), ( |
| 184 | f"{context}: missing an unambiguously top-{requested} token" |
| 185 | ) |
| 186 | assert (reference[top_ids] >= boundary - 2 * delta).all(), ( |
| 187 | f"{context}: returned token below the top-{requested} boundary" |
| 188 | ) |
| 189 | return errors |
| 190 | |
| 191 | |
| 192 | def check_greedy_token( |
| 193 | logits: torch.Tensor, |
| 194 | token: int, |
| 195 | tolerance: AccuracyTolerance, |
| 196 | context: str, |
| 197 | ) -> float: |
| 198 | assert not torch.isnan(logits).any(), context |
| 199 | assert torch.isfinite(logits[token]), ( |
| 200 | f"{context}: selected token {token} is forbidden by the sampling parameters" |
| 201 | ) |
| 202 | best = int(logits.argmax()) |
| 203 | gap = float(logits[best] - logits[token]) |
| 204 | assert gap <= tolerance.greedy_atol, ( |
| 205 | f"{context}: selected token={token}, reference best={best}, " |
| 206 | f"logit gap={gap}, bound={tolerance.greedy_atol}" |
| 207 | ) |
| 208 | return gap |
| 209 | |
| 210 | |
| 211 | def check_accuracy_budget(errors, gaps, tolerance, context): |
| 212 | """Reject systematic sub-bound errors within each request and score stream.""" |
| 213 | if errors: |
| 214 | finite_errors = torch.cat(errors) |
| 215 | if finite_errors.numel(): |
| 216 | mean_error = float(finite_errors.mean()) |
| 217 | assert mean_error <= tolerance.logprob_mean_atol, ( |
| 218 | f"{context}: mean logprob error={mean_error}, " |
| 219 | f"bound={tolerance.logprob_mean_atol}" |
| 220 | ) |
| 221 | assert sum(gaps) <= tolerance.greedy_total_atol, ( |
| 222 | f"{context}: total greedy logit gap={sum(gaps)}, " |
| 223 | f"bound={tolerance.greedy_total_atol}" |
| 224 | ) |
| 225 | |
| 226 | |
| 227 | def _check_decoded_tokens(processor, scores, history, context): |
| 228 | ids = list(scores) |
| 229 | decoded = convert_ids_list_to_tokens(processor.tokenizer, ids) |
| 230 | decoded = processor._verify_tokens(decoded, ids, context_token_ids=history[-4:]) |
| 231 | for token, text in zip(ids, decoded, strict=True): |
| 232 | assert scores[token].decoded_token == text, ( |
| 233 | f"{context}: incorrect decoded text for token {token}" |
| 234 | ) |
| 235 | |
| 236 | |
| 237 | @torch.inference_mode() |
| 238 | def check_request_accuracy( |
| 239 | hf: HfRunner, |
| 240 | request: RequestOutput, |
| 241 | params: SamplingParams, |
| 242 | tolerance: AccuracyTolerance, |
| 243 | context: str, |
| 244 | ) -> None: |
| 245 | """Score the actual continuation, including positions after any divergence.""" |
| 246 | assert request.finished and len(request.outputs) == 1, context |
| 247 | (completion,) = request.outputs |
| 248 | prompt = request.prompt_token_ids |
| 249 | assert prompt is not None and len(prompt) > 0, context |
| 250 | tokens = list(completion.token_ids) |
| 251 | assert params.max_tokens is not None and 0 < len(tokens) <= params.max_tokens, ( |
| 252 | context |
| 253 | ) |
| 254 | # Incremental decoding buffers incomplete UTF-8; decoding the final token |
| 255 | # list in one call can incorrectly append a replacement character. |
| 256 | stream = DecodeStream(skip_special_tokens=params.skip_special_tokens) |
| 257 | tokenizer = hf.tokenizer.backend_tokenizer |
| 258 | for token in prompt: |
| 259 | stream.step(tokenizer, token) |
| 260 | visible = tokens |
| 261 | if completion.finish_reason == "stop" and not params.include_stop_str_in_output: |
| 262 | visible = tokens[:-1] |
| 263 | expected_text = "".join(stream.step(tokenizer, token) or "" for token in visible) |
| 264 | assert completion.text == expected_text, ( |
| 265 | f"{context}: decoded text does not match generated token IDs" |
| 266 | ) |
| 267 | inputs = torch.tensor([prompt + tokens], device=hf.device) |
| 268 | logits = hf.model(input_ids=inputs, use_cache=False).logits[0].float() |
| 269 | # Use the existing byte-fallback handling with an independent request |
| 270 | # history; numerical scores always come from the HF model above. |
| 271 | decoder = LogprobsProcessor( |
| 272 | tokenizer=hf.tokenizer, |
| 273 | logprobs=None, |
| 274 | prompt_logprobs=None, |
| 275 | cumulative_logprob=None, |
| 276 | num_logprobs=None, |
| 277 | num_prompt_logprobs=None, |
| 278 | ) |
| 279 | # Prompt token i is predicted by the causal logits at position i - 1. |
| 280 | if params.prompt_logprobs is None: |
| 281 | assert request.prompt_logprobs is None, context |
| 282 | else: |
| 283 | assert request.prompt_logprobs is not None, context |
| 284 | assert len(request.prompt_logprobs) == len(prompt), context |
| 285 | assert request.prompt_logprobs[0] is None, context |
| 286 | errors = [] |
| 287 | for i, scores in enumerate(request.prompt_logprobs[1:], 1): |
| 288 | assert scores is not None, context |
| 289 | _check_decoded_tokens(decoder, scores, prompt[1:i], context) |
| 290 | errors.append( |
| 291 | check_logprobs( |
| 292 | scores, |
| 293 | logits[i - 1].log_softmax(-1), |
| 294 | prompt[i], |
| 295 | params.prompt_logprobs, |
| 296 | tolerance, |
| 297 | f"{context}, prompt position={i}", |
| 298 | ) |
| 299 | ) |
| 300 | check_accuracy_budget(errors, [], tolerance, f"{context}, prompt logprobs") |
| 301 | |
| 302 | if params.logprobs is None: |
| 303 | assert completion.logprobs is None, context |
| 304 | else: |
| 305 | assert completion.logprobs is not None, context |
| 306 | assert len(completion.logprobs) == len(tokens), context |
| 307 | |
| 308 | reference = SamplingReference( |
| 309 | hf.tokenizer, |
| 310 | logits.shape[-1], |
| 311 | hf.model.generation_config.to_dict(), |
| 312 | params, |
| 313 | ) |
| 314 | errors = [] |
| 315 | gaps = [] |
| 316 | for i, token in enumerate(tokens): |
| 317 | raw, processed = reference.distributions( |
| 318 | logits[len(prompt) + i - 1], tokens[:i] |
| 319 | ) |
| 320 | position = f"{context}, generated position={i}" |
| 321 | gaps.append(check_greedy_token(processed, token, tolerance, position)) |
| 322 | if completion.logprobs is not None: |
| 323 | _check_decoded_tokens(decoder, completion.logprobs[i], tokens[:i], position) |
| 324 | errors.append( |
| 325 | check_logprobs( |
| 326 | completion.logprobs[i], |
| 327 | raw.log_softmax(-1), |
| 328 | token, |
| 329 | params.logprobs, |
| 330 | tolerance, |
| 331 | position, |
| 332 | ) |
| 333 | ) |
| 334 | check_accuracy_budget(errors, gaps, tolerance, f"{context}, completion") |
| 335 | # Advance the grammar through the last token too; an incomplete JSON prefix |
| 336 | # is allowed at max_tokens, but every emitted token must remain valid. |
| 337 | reference.accept_history(tokens) |
| 338 | check_stopping( |
| 339 | tokens, |
| 340 | completion.finish_reason, |
| 341 | reference.stop_ids, |
| 342 | params, |
| 343 | context, |
| 344 | eos_token_id=hf.tokenizer.eos_token_id, |
| 345 | stop_reason=completion.stop_reason, |
| 346 | ) |
| 347 | |
| 348 | |
| 349 | def check_stopping( |
| 350 | tokens, |
| 351 | finish_reason, |
| 352 | stop_ids, |
| 353 | params, |
| 354 | context, |
| 355 | *, |
| 356 | eos_token_id: int | None, |
| 357 | stop_reason: int | str | None, |
| 358 | ): |
| 359 | assert not params.ignore_eos and not params.stop, context |
| 360 | assert not stop_ids.intersection(tokens[:-1]), ( |
| 361 | f"{context}: generation continued after a stop token" |
| 362 | ) |
| 363 | if tokens[-1] in stop_ids: |
| 364 | assert finish_reason == "stop", ( |
| 365 | f"{context}: stop token did not terminate output" |
| 366 | ) |
| 367 | expected_stop_reason = None if tokens[-1] == eos_token_id else tokens[-1] |
| 368 | assert stop_reason == expected_stop_reason, ( |
| 369 | f"{context}: stop reason={stop_reason}, expected={expected_stop_reason}" |
| 370 | ) |
| 371 | else: |
| 372 | assert finish_reason == "length" and len(tokens) == params.max_tokens, context |
| 373 | assert stop_reason is None, f"{context}: length cap has a stop reason" |
| 374 | |
| 375 | |
| 376 | @lru_cache(maxsize=4) |
| 377 | def _compiler(tokenizer, vocab_size: int): |
| 378 | info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size) |
| 379 | return xgr.GrammarCompiler(info, max_threads=1, cache_enabled=True) |
| 380 | |
| 381 | |
| 382 | class SamplingReference: |
| 383 | """One request's reference; histories must grow without changing their prefix.""" |
| 384 | |
| 385 | def __init__( |
| 386 | self, |
| 387 | tokenizer, |
| 388 | vocab_size: int, |
| 389 | generation_config: dict, |
| 390 | params: "SamplingParams", |
| 391 | ): |
| 392 | assert params.temperature == 0.0, "Only the original greedy matrix is supported" |
| 393 | assert params.presence_penalty == 0.0 and params.repetition_penalty == 1.0 |
| 394 | assert params.allowed_token_ids is None and not params.logit_bias |
| 395 | self.vocab_size = vocab_size |
| 396 | self.min_tokens = params.min_tokens |
| 397 | self.frequency_penalty = params.frequency_penalty |
| 398 | self.stop_ids = set(params.stop_token_ids or []) |
| 399 | for eos in (tokenizer.eos_token_id, generation_config.get("eos_token_id")): |
| 400 | if eos is not None: |
| 401 | self.stop_ids.update([eos] if isinstance(eos, int) else eos) |
| 402 | assert all(0 <= token < vocab_size for token in self.stop_ids) |
| 403 | |
| 404 | # Match public bad-word tokenization semantics without using sampler code. |
| 405 | self.bad_words: list[list[int]] = [] |
| 406 | for word in params.bad_words or []: |
| 407 | plain = tokenizer.encode(word.lstrip(), add_special_tokens=False) |
| 408 | spaced = tokenizer.encode(" " + word.lstrip(), add_special_tokens=False) |
| 409 | assert plain, f"Empty bad-word tokenization: {word!r}" |
| 410 | self.bad_words.append(plain) |
| 411 | if spaced and len(spaced) == len(plain) and spaced[0] != plain[0]: |
| 412 | self.bad_words.append(spaced) |
| 413 | |
| 414 | self.history: list[int] = [] |
| 415 | self.matcher = None |
| 416 | self.bitmask = None |
| 417 | if params.structured_outputs is not None: |
| 418 | schema = params.structured_outputs.json |
| 419 | assert schema is not None, "Only the original JSON schema is supported" |
| 420 | # Preserve dictionary order: xgrammar's default grammar orders fields. |
| 421 | schema = schema if isinstance(schema, str) else json.dumps(schema) |
| 422 | compiled = _compiler(tokenizer, vocab_size).compile_json_schema( |
| 423 | schema, any_whitespace=True |
| 424 | ) |
| 425 | self.matcher = xgr.GrammarMatcher( |
| 426 | compiled, override_stop_tokens=sorted(self.stop_ids) or None |
| 427 | ) |
| 428 | self.bitmask = xgr.allocate_token_bitmask(1, vocab_size) |
| 429 | |
| 430 | def accept_history(self, generated_ids: list[int]) -> bool: |
| 431 | """Validate newly generated grammar tokens, including the final EOS. |
| 432 | |
| 433 | Return whether the grammar has terminated. Use this after the final |
| 434 | output token; no next-token distribution exists after termination. |
| 435 | """ |
| 436 | assert generated_ids[: len(self.history)] == self.history, ( |
| 437 | "Each reference belongs to one monotonically growing request history" |
| 438 | ) |
| 439 | if self.matcher is not None: |
| 440 | for token in generated_ids[len(self.history) :]: |
| 441 | assert self.matcher.accept_token(token), ( |
| 442 | f"Grammar rejected token {token} at position {len(self.history)}" |
| 443 | ) |
| 444 | self.history.append(token) |
| 445 | return self.matcher.is_terminated() |
| 446 | else: |
| 447 | self.history = generated_ids.copy() |
| 448 | return False |
| 449 | |
| 450 | def distributions( |
| 451 | self, raw_logits: torch.Tensor, generated_ids: list[int] |
| 452 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 453 | """Return FP32 logits before and after non-grammar sampling transforms. |
| 454 | |
| 455 | ``raw_logits`` scores the next token after the prompt and |
| 456 | ``generated_ids``. Returned tensors do not alias the model's tensor. |
| 457 | Normalize the first result to check reported completion logprobs; |
| 458 | the second supplies the greedy decision reference. |
| 459 | """ |
| 460 | assert raw_logits.shape == (self.vocab_size,) |
| 461 | assert not self.accept_history(generated_ids), ( |
| 462 | "Cannot generate a distribution after grammar EOS" |
| 463 | ) |
| 464 | |
| 465 | grammar_logits = raw_logits.float().clone() |
| 466 | if self.matcher is not None: |
| 467 | assert self.bitmask is not None |
| 468 | self.matcher.fill_next_token_bitmask(self.bitmask) |
| 469 | words = self.bitmask[0, :, None] |
| 470 | allowed = ((words >> torch.arange(32)) & 1).flatten() |
| 471 | allowed = allowed[: self.vocab_size].bool().to(raw_logits.device) |
| 472 | grammar_logits.masked_fill_(~allowed, -torch.inf) |
| 473 | |
| 474 | processed = grammar_logits.clone() |
| 475 | for bad_word in self.bad_words: |
| 476 | prefix = bad_word[:-1] |
| 477 | if not prefix or generated_ids[-len(prefix) :] == prefix: |
| 478 | processed[bad_word[-1]] = -torch.inf |
| 479 | |
| 480 | if len(generated_ids) < self.min_tokens and self.stop_ids: |
| 481 | stop_ids = sorted(self.stop_ids) |
| 482 | stop_scores = processed[stop_ids].clone() |
| 483 | processed[stop_ids] = -torch.inf |
| 484 | # A completed grammar may have no continuation except a stop token. |
| 485 | if self.matcher is not None and torch.isneginf(processed).all(): |
| 486 | processed[stop_ids] = stop_scores |
| 487 | |
| 488 | if self.frequency_penalty: |
| 489 | for token, count in Counter(generated_ids).items(): |
| 490 | processed[token] -= self.frequency_penalty * count |
| 491 | return grammar_logits, processed |