@tduncan / Nvidia-Optimization-Stack-VLLM / commits / df8fd42116

[CI][Spec Decode] Add async scheduling accuracy tests (#55840)

Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Signed-off-by: Andreas Karatzas <andreas.karatzas@protonmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Andreas Karatzas committed Sep 30, 2026 at 20:06 UTC df8fd42116f172b7a53bc10c8a680b05232edbed
4 files changed +1166 -3
.buildkite/test-amd.yaml
+29 -1
@@ -454,7 +454,7 @@ steps:
454 - tests/v1/e2e/general/
455 - vllm/platforms/rocm.py
456 commands:
457 - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
457 + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_async_scheduling_accuracy.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
458
459 - label: ":amd: (MI250) E2E Scheduling" # TBD
460 timeout_in_minutes: 55
@@ -470,6 +470,20 @@ steps:
470 commands:
471 - pytest -v -s v1/e2e/general/test_async_scheduling.py
472
473 +- label: ":amd: (MI250) E2E Scheduling Accuracy" # TBD
474 + timeout_in_minutes: 55
475 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
476 + dind: false
477 + agent_pool: mi250_1
478 + optional: true
479 + working_dir: "/vllm-workspace/tests"
480 + source_file_dependencies:
481 + - vllm/v1/
482 + - tests/v1/e2e/general/
483 + - vllm/platforms/rocm.py
484 + commands:
485 + - pytest -v -s v1/e2e/general/test_async_scheduling_accuracy.py
486 +
487 - label: ":amd: (MI250) V1 Engine" # TBD
488 timeout_in_minutes: 45
489 mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -2654,6 +2668,20 @@ steps:
2668 commands:
2669 - pytest -v -s v1/e2e/general/test_async_scheduling.py
2670
2671 +- label: ":amd: (MI300) E2E Scheduling Accuracy" # TBD
2672 + timeout_in_minutes: 55
2673 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
2674 + dind: false
2675 + agent_pool: mi300_1
2676 + optional: true
2677 + working_dir: "/vllm-workspace/tests"
2678 + source_file_dependencies:
2679 + - vllm/v1/
2680 + - tests/v1/e2e/general/
2681 + - vllm/platforms/rocm.py
2682 + commands:
2683 + - pytest -v -s v1/e2e/general/test_async_scheduling_accuracy.py
2684 +
2685 - label: ":amd: (MI300) V1 Engine" # TBD
2686 timeout_in_minutes: 45
2687 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
.buildkite/test_areas/engine.yaml
+22 -2
@@ -85,6 +85,26 @@ steps:
85 source_file_dependencies:
86 - vllm/platforms/rocm.py
87
88 +- label: ":nvidia: (H200 MIG 18GB) E2E Scheduling Accuracy"
89 + key: e2e-scheduling-accuracy-1-gpu
90 + timeout_in_minutes: 53
91 + device: h200_18gb
92 + source_file_dependencies:
93 + - vllm/v1/
94 + - tests/v1/e2e/general/
95 + commands:
96 + - pytest -v -s v1/e2e/general/test_async_scheduling_accuracy.py
97 + mirror:
98 + amd:
99 + label: ":amd: (MI250) E2E Scheduling Accuracy"
100 + dind: false
101 + device: mi250_1
102 + timeout_in_minutes: 55
103 + depends_on:
104 + - image-build-amd
105 + source_file_dependencies:
106 + - vllm/platforms/rocm.py
107 +
108 - label: ":nvidia: (H200 MIG 18GB) E2E Core"
109 device: h200_18gb
110 key: e2e-core-1-gpu
@@ -93,13 +113,13 @@ steps:
113 - vllm/v1/
114 - tests/v1/e2e/general/
115 commands:
96 - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
116 + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_async_scheduling_accuracy.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
117 mirror:
118 amd:
119 label: ":amd: (MI250) E2E Core"
120 dind: false
121 commands:
102 - - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
122 + - pytest -v -s v1/e2e/general --ignore v1/e2e/general/test_async_scheduling.py --ignore v1/e2e/general/test_async_scheduling_accuracy.py --ignore v1/e2e/general/test_kv_sharing_fast_prefill.py -k "not test_mamba_prefix_cache_mrv2"
123 device: mi250_1
124 timeout_in_minutes: 50
125 depends_on:
tests/v1/e2e/general/accuracy_utils.py new
+491
@@ -0,0 +1,491 @@
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
tests/v1/e2e/general/test_async_scheduling_accuracy.py new
+624
@@ -0,0 +1,624 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Accuracy with automatic attention selection and the model's normal dtype.
4 +
5 +Complements the strict Flex/FP32 invariance checks in test_async_scheduling.py.
6 +"""
7 +
8 +import os
9 +from dataclasses import replace
10 +from typing import Any
11 +
12 +import pytest
13 +import torch
14 +
15 +from tests.conftest import HfRunner, VllmRunner
16 +from tests.utils import (
17 + large_gpu_mark,
18 + single_gpu_only,
19 +)
20 +from tests.v1.e2e.general.accuracy_utils import (
21 + AccuracyTolerance,
22 + check_accuracy_budget,
23 + check_first_divergence,
24 + check_greedy_token,
25 + check_logprobs,
26 + check_request_accuracy,
27 + check_stopping,
28 +)
29 +from vllm import SamplingParams
30 +from vllm.logprobs import Logprob
31 +from vllm.outputs import CompletionOutput, RequestOutput
32 +from vllm.platforms import current_platform
33 +from vllm.sampling_params import StructuredOutputsParams
34 +from vllm.utils.torch_utils import set_default_torch_num_threads
35 +from vllm.v1.metrics.reader import Metric
36 +
37 +MODEL = "Qwen/Qwen3-0.6B"
38 +EAGLE3_TARGET_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
39 +
40 +# Need to enforce eager for MRV2 while we sort out cudagraph issues.
41 +ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1"
42 +
43 +first_prompt = (
44 + "The following numbers of the sequence "
45 + + ", ".join(str(i) for i in range(10))
46 + + " are:"
47 +)
48 +example_prompts = [first_prompt, "In one word, the capital of France is "] + [
49 + f"Tell me about the number {i}: " for i in range(32)
50 +]
51 +
52 +default_params = dict(
53 + temperature=0.0, # greedy
54 + max_tokens=30,
55 + min_tokens=28,
56 +)
57 +
58 +
59 +@single_gpu_only
60 +def test_without_spec_decoding(
61 + vllm_runner: type[VllmRunner],
62 + hf_runner: type[HfRunner],
63 + sample_json_schema,
64 +):
65 + """Check target-model accuracy across scheduling and executor configurations."""
66 + struct_outputs = StructuredOutputsParams(json=sample_json_schema)
67 + test_sampling_params: list[dict[str, Any]] = [
68 + dict(),
69 + dict(frequency_penalty=-1.0),
70 + dict(bad_words=["the", " the"]),
71 + dict(bad_words=["the", " the"], logprobs=2),
72 + dict(logprobs=2),
73 + dict(logprobs=2, frequency_penalty=-1.0),
74 + dict(prompt_logprobs=2),
75 + dict(prompt_logprobs=2, logprobs=2),
76 + dict(structured_outputs=struct_outputs),
77 + dict(
78 + structured_outputs=struct_outputs,
79 + logprobs=2,
80 + ),
81 + dict(
82 + structured_outputs=struct_outputs,
83 + frequency_penalty=-1.0,
84 + ),
85 + dict(
86 + structured_outputs=struct_outputs,
87 + logprobs=2,
88 + frequency_penalty=-1.0,
89 + ),
90 + ]
91 +
92 + # test_preemption, executor, async_scheduling,
93 + # spec_config, test_prefill_chunking
94 + test_configs = [
95 + (False, "mp", False, None, False),
96 + (True, "mp", False, None, True),
97 + (False, "mp", True, None, False),
98 + (False, "uni", True, None, False),
99 + (True, "mp", True, None, False),
100 + (True, "uni", True, None, False),
101 + (False, "mp", True, None, True),
102 + (True, "mp", True, None, True),
103 + (True, "uni", True, None, True),
104 + ]
105 +
106 + run_tests(vllm_runner, hf_runner, MODEL, test_configs, test_sampling_params)
107 +
108 +
109 +@single_gpu_only
110 +@large_gpu_mark(min_gb=16)
111 +def test_with_eagle3_spec_decoding(
112 + vllm_runner: type[VllmRunner],
113 + hf_runner: type[HfRunner],
114 + sample_json_schema,
115 +):
116 + """Check accuracy and acceptance, including drafts that exceed model length."""
117 + spec_config = {
118 + "method": "eagle3",
119 + "num_speculative_tokens": 2,
120 + "model": "nm-testing/Llama3_2_1B_speculator.eagle3",
121 + }
122 + # Set small draft model len to force doesn't-fit-in-drafter case.
123 + spec_config_short = spec_config | {"max_model_len": 50}
124 +
125 + struct_outputs = StructuredOutputsParams(json=sample_json_schema)
126 +
127 + test_sampling_params: list[dict[str, Any]] = [
128 + dict(),
129 + dict(frequency_penalty=-1.0),
130 + dict(bad_words=["the", " the"]),
131 + dict(bad_words=["the", " the"], logprobs=2),
132 + dict(logprobs=2),
133 + dict(logprobs=2, frequency_penalty=-1.0),
134 + dict(prompt_logprobs=2),
135 + dict(prompt_logprobs=2, logprobs=2),
136 + dict(structured_outputs=struct_outputs),
137 + dict(
138 + structured_outputs=struct_outputs,
139 + logprobs=2,
140 + frequency_penalty=-1.0,
141 + ),
142 + ]
143 +
144 + # test_preemption, executor, async_scheduling,
145 + # spec_config, test_prefill_chunking
146 + test_configs = [
147 + (False, "mp", False, None, False),
148 + (False, "mp", False, spec_config, False),
149 + (True, "mp", False, spec_config, True),
150 + (True, "uni", False, spec_config_short, True),
151 + (False, "mp", True, spec_config, False),
152 + (True, "mp", True, spec_config, False),
153 + (False, "mp", True, spec_config_short, True),
154 + (True, "uni", True, spec_config, False),
155 + (True, "uni", True, spec_config_short, False),
156 + (True, "mp", True, spec_config, True),
157 + (True, "uni", True, spec_config_short, True),
158 + ]
159 +
160 + run_tests(
161 + vllm_runner, hf_runner, EAGLE3_TARGET_MODEL, test_configs, test_sampling_params
162 + )
163 +
164 +
165 +@single_gpu_only
166 +@pytest.mark.skipif(
167 + current_platform.is_xpu(),
168 + reason=("XPU matmul/attention kernels are not batch-invariant"),
169 +)
170 +def test_with_ngram_gpu_spec_decoding(
171 + vllm_runner: type[VllmRunner], hf_runner: type[HfRunner]
172 +):
173 + """Check ngram accuracy and acceptance across scheduling configurations."""
174 + ngram_gpu_config = {
175 + "method": "ngram_gpu",
176 + "num_speculative_tokens": 3,
177 + "prompt_lookup_max": 3,
178 + "prompt_lookup_min": 2,
179 + }
180 +
181 + # test_preemption, executor, async_scheduling,
182 + # spec_config, test_prefill_chunking
183 + test_configs = [
184 + (False, "mp", False, None, False),
185 + (False, "mp", False, ngram_gpu_config, False),
186 + (True, "mp", False, ngram_gpu_config, True),
187 + (False, "mp", True, ngram_gpu_config, False),
188 + (True, "mp", True, ngram_gpu_config, False),
189 + (True, "uni", True, ngram_gpu_config, False),
190 + (True, "mp", True, ngram_gpu_config, True),
191 + ]
192 +
193 + run_tests(vllm_runner, hf_runner, MODEL, test_configs, [{}, dict(logprobs=2)])
194 +
195 +
196 +def run_tests(
197 + vllm_runner: type[VllmRunner],
198 + hf_runner: type[HfRunner],
199 + model: str,
200 + test_configs: list[tuple],
201 + test_sampling_params: list[dict[str, Any]],
202 +):
203 + """Compare the shared prefix and independently validate every continuation."""
204 + assert test_configs[0] == (False, "mp", False, None, False), (
205 + "The baseline must use synchronous target-only decoding"
206 + )
207 + outputs = []
208 + for n, config in enumerate(test_configs, 1):
209 + outputs.append(
210 + run_test(
211 + vllm_runner,
212 + model,
213 + f"{n}/{len(test_configs)}",
214 + test_sampling_params,
215 + *config,
216 + )
217 + )
218 +
219 + baseline_config, baseline_batches, _ = outputs[0]
220 + # Compare vLLM's own decisions at the first divergence, where histories
221 + # still match. HF then scores each continuation, including after divergence.
222 + # Neither engine forces an attention backend, dtype, or matmul precision.
223 + # Qwen/Llama BF16 controls and Qwen holdouts measured max score error <0.38,
224 + # per-stream request mean <0.073, max greedy gap 0.125, and request sum 0.25.
225 + # Per-request budgets also reject systematic errors below the scalar bounds.
226 + tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
227 + # These short CPU reductions cost more to distribute across a thread pool.
228 + with set_default_torch_num_threads(1), hf_runner(model) as hf:
229 + hf.model.eval()
230 + for config, batches, _ in outputs:
231 + assert len(batches) == len(test_sampling_params)
232 + for baseline_batch, batch, overrides in zip(
233 + baseline_batches, batches, test_sampling_params, strict=True
234 + ):
235 + assert len(batch) == len(example_prompts), config
236 + params = SamplingParams(**default_params, **overrides)
237 + for i, (baseline, request, prompt) in enumerate(
238 + zip(baseline_batch, batch, example_prompts, strict=True)
239 + ):
240 + assert request.prompt_token_ids == hf.tokenizer.encode(prompt)
241 + context = f"config=[{config}], params={overrides}, request={i}"
242 + if params.logprobs is not None:
243 + check_first_divergence(
244 + baseline,
245 + request,
246 + params,
247 + tolerance,
248 + f"baseline=[{baseline_config}], {context}",
249 + )
250 + check_request_accuracy(
251 + hf,
252 + request,
253 + params,
254 + tolerance,
255 + context,
256 + )
257 + print(f"ACCURACY PASSED: config=[{config}], params={overrides}")
258 +
259 + baseline_acceptances = next((o[2] for o in outputs if o[2] is not None), None)
260 + if baseline_acceptances is not None:
261 + for settings, (config, _, acceptances) in zip(
262 + test_configs, outputs, strict=True
263 + ):
264 + if acceptances is None:
265 + continue
266 + test_preemption, _, _, spec_config, _ = settings
267 + assert spec_config is not None
268 + for baseline, actual, params in zip(
269 + baseline_acceptances, acceptances, test_sampling_params, strict=True
270 + ):
271 + context = f"config=[{config}], params={params}"
272 + if spec_config.get("max_model_len") is None:
273 + # Keep the original acceptance-quality floor per batch.
274 + relative_drop = (
275 + 0.10 if current_platform.is_rocm() and test_preemption else 0.05
276 + )
277 + assert actual >= baseline * (1 - relative_drop), (
278 + f"{context}: acceptance={actual}, baseline={baseline}"
279 + )
280 + else:
281 + assert actual > 0.1, f"{context}: acceptance={actual}"
282 +
283 +
284 +def run_test(
285 + vllm_runner: type[VllmRunner],
286 + model: str,
287 + test_str: str,
288 + sampling_param_tests: list[dict[str, Any]],
289 + test_preemption: bool,
290 + executor: str,
291 + async_scheduling: bool,
292 + spec_config: dict[str, Any] | None,
293 + test_prefill_chunking: bool,
294 +):
295 + spec_decoding = spec_config is not None
296 + spec_method = (spec_config or {}).get("method", "none")
297 + # Chunked ngram decoding admits fewer concurrent requests under the
298 + # 48-token budget. Its original 33-block cache did not preempt on ROCm.
299 + # A 17-block cache forces contention while every original request fits
300 + # comfortably within 256 tokens (the longest prompt plus output is 67).
301 + cache_blocks = 17 if test_prefill_chunking and spec_method == "ngram_gpu" else 33
302 + cache_arg: dict[str, Any] = (
303 + dict(
304 + num_gpu_blocks_override=cache_blocks,
305 + max_model_len=(cache_blocks - 1) * 16,
306 + )
307 + if test_preemption
308 + else dict(gpu_memory_utilization=0.9, max_model_len=4096)
309 + )
310 + spec_mml = (spec_config or {}).get("max_model_len")
311 + test_config = (
312 + f"executor={executor}, preemption={test_preemption}, "
313 + f"async_sched={async_scheduling}, "
314 + f"chunk_prefill={test_prefill_chunking}, "
315 + f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}"
316 + )
317 + print("-" * 80)
318 + print(f"---- TESTING {test_str}: {test_config}")
319 + print("-" * 80)
320 +
321 + with vllm_runner(
322 + model,
323 + enable_chunked_prefill=test_prefill_chunking,
324 + # Force prefill chunking
325 + max_num_batched_tokens=48 if test_prefill_chunking else None,
326 + enforce_eager=ENFORCE_EAGER,
327 + async_scheduling=async_scheduling,
328 + distributed_executor_backend=executor,
329 + speculative_config=spec_config,
330 + disable_log_stats=False,
331 + enable_prefix_caching=False if current_platform.is_rocm() else None,
332 + **cache_arg,
333 + ) as vllm_model:
334 + results = []
335 + acceptance_rates: list[float] | None = [] if spec_decoding else None
336 + for override_params in sampling_param_tests:
337 + metrics_before = vllm_model.llm.get_metrics()
338 + print(f"----------- RUNNING PARAMS: {override_params}")
339 + results.append(
340 + vllm_model.llm.generate(
341 + vllm_model.get_inputs(example_prompts),
342 + sampling_params=SamplingParams(**default_params, **override_params),
343 + )
344 + )
345 + metrics_after = vllm_model.llm.get_metrics()
346 + if acceptance_rates is not None:
347 + acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
348 + acceptance_rates.append(acceptance_rate)
349 + print(f"ACCEPTANCE RATE {acceptance_rate}")
350 +
351 + if test_preemption:
352 + preemptions = _get_count(
353 + metrics_before, metrics_after, "vllm:num_preemptions"
354 + )
355 + assert preemptions > 0, "preemption test had no preemptions"
356 +
357 + # Preserve the original parameter-effect checks across repeated batches.
358 + if len(results) > 1:
359 + baseline = _result_signature(results[0])
360 + for outputs, params in zip(results[1:], sampling_param_tests[1:], strict=True):
361 + assert _result_signature(outputs) != baseline, (
362 + f"{test_config}: sampling parameters had no observable effect: {params}"
363 + )
364 + # The old ROCm filter used schema-only as its first parameter set.
365 + # Retain those comparisons after restoring the ordinary cases too.
366 + schema_batches = [
367 + batch
368 + for batch, params in zip(results, sampling_param_tests, strict=True)
369 + if params.get("structured_outputs") is not None
370 + ]
371 + for batch in schema_batches[1:]:
372 + assert _result_signature(batch) != _result_signature(schema_batches[0]), (
373 + f"{test_config}: structured sampling parameters had no effect"
374 + )
375 +
376 + return test_config, results, acceptance_rates
377 +
378 +
379 +def _result_signature(requests):
380 + return [
381 + (
382 + list(request.outputs[0].token_ids),
383 + request.outputs[0].text,
384 + request.prompt_logprobs,
385 + request.outputs[0].logprobs,
386 + )
387 + for request in requests
388 + ]
389 +
390 +
391 +def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
392 + draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
393 + accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
394 + assert draft > 0, "speculative test did not draft any tokens"
395 + assert 0 <= accept <= draft, f"invalid acceptance counters: {accept=}, {draft=}"
396 + return accept / draft
397 +
398 +
399 +def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
400 + before_val = next(m.value for m in before if m.name == name)
401 + after_val = next(m.value for m in after if m.name == name)
402 + return after_val - before_val
403 +
404 +
405 +@pytest.mark.parametrize(
406 + "corruption", ["score", "rank", "top_k", "missing", "nan", "positive", "order"]
407 +)
408 +def test_accuracy_checks_reject_incorrect_logprobs(corruption):
409 + second = 2.875 if corruption == "order" else 1.0
410 + reference = torch.tensor([3.0, second, -2.0, -5.0]).log_softmax(-1)
411 + tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
412 + scores = {
413 + 0: Logprob(logprob=float(reference[0]), rank=1, decoded_token="a"),
414 + 1: Logprob(logprob=float(reference[1]), rank=2, decoded_token="b"),
415 + }
416 + check_logprobs(scores, reference, 0, 2, tolerance, "control")
417 + if corruption == "score":
418 + scores[0] = replace(scores[0], logprob=scores[0].logprob - 1.0)
419 + elif corruption == "rank":
420 + scores[1] = replace(scores[1], rank=1)
421 + elif corruption == "top_k":
422 + del scores[1]
423 + scores[3] = Logprob(logprob=float(reference[3]), rank=2, decoded_token="d")
424 + elif corruption == "missing":
425 + del scores[0]
426 + elif corruption == "positive":
427 + scores[0] = replace(scores[0], logprob=0.1)
428 + elif corruption == "order":
429 + scores[0] = replace(scores[0], rank=2)
430 + scores[1] = replace(scores[1], rank=1)
431 + else:
432 + scores[0] = replace(scores[0], logprob=float("nan"))
433 + with pytest.raises(AssertionError):
434 + check_logprobs(scores, reference, 0, 2, tolerance, corruption)
435 +
436 +
437 +def test_accuracy_checks_distinguish_ties_from_wrong_tokens():
438 + tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
439 + logits = torch.tensor([3.0, 2.875, -2.0, -torch.inf])
440 + check_greedy_token(logits, 0, tolerance, "best")
441 + check_greedy_token(logits, 1, tolerance, "near tie")
442 + for token in (2, 3):
443 + with pytest.raises(AssertionError):
444 + check_greedy_token(logits, token, tolerance, "incorrect token")
445 +
446 + # top-k may omit a sampled token tied at its boundary. The sampled rank
447 + # can duplicate a sequential top-k rank without corrupting the result.
448 + reference = torch.zeros(4).log_softmax(-1)
449 + scores = {
450 + token: Logprob(logprob=float(reference[token]), rank=rank)
451 + for token, rank in [(3, 1), (0, 1), (1, 2)]
452 + }
453 + check_logprobs(scores, reference, 3, 2, tolerance, "tied top-k")
454 +
455 + # A grammar with one allowed token still returns the requested top-k,
456 + # including masked entries whose score is exactly negative infinity.
457 + reference = torch.tensor([0.0, -torch.inf, -torch.inf])
458 + scores = {0: Logprob(0.0, 1), 1: Logprob(-float("inf"), 2)}
459 + check_logprobs(scores, reference, 0, 2, tolerance, "masked top-k")
460 +
461 +
462 +def test_accuracy_checks_reject_incorrect_stopping():
463 + params = SamplingParams(max_tokens=3)
464 + stop_ids = {2, 3}
465 + for tokens, reason, stop_reason in [
466 + ([0, 1, 2], "stop", None),
467 + ([0, 1, 3], "stop", 3),
468 + ([0, 1, 0], "length", None),
469 + ]:
470 + check_stopping(
471 + tokens,
472 + reason,
473 + stop_ids,
474 + params,
475 + "valid stopping at length cap",
476 + eos_token_id=2,
477 + stop_reason=stop_reason,
478 + )
479 + for tokens, reason, stop_reason in [
480 + ([0, 2, 1], "length", None),
481 + ([0, 1, 2], "length", None),
482 + ([0, 1, 2], "stop", 2),
483 + ([0, 1, 3], "stop", None),
484 + ([0, 1, 3], "stop", 2),
485 + ([0, 1, 0], "length", 3),
486 + ]:
487 + with pytest.raises(AssertionError):
488 + check_stopping(
489 + tokens,
490 + reason,
491 + stop_ids,
492 + params,
493 + "incorrect stopping",
494 + eos_token_id=2,
495 + stop_reason=stop_reason,
496 + )
497 +
498 +
499 +def test_accuracy_checks_reject_systematic_sub_bound_errors():
500 + tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
501 + reference = torch.tensor([3.0, 2.75, -2.0]).log_softmax(-1)
502 + scores = {i: Logprob(float(reference[i]) - 0.4, i + 1) for i in (0, 1)}
503 + # Each corrupted score fits the scalar bound, but the bias must not pass.
504 + errors = check_logprobs(scores, reference, 0, 2, tolerance, "biased scores")
505 + with pytest.raises(AssertionError, match="mean logprob error"):
506 + check_accuracy_budget([errors], [], tolerance, "biased scores")
507 +
508 + gap = check_greedy_token(reference, 1, tolerance, "one ambiguous choice")
509 + with pytest.raises(AssertionError, match="total greedy logit gap"):
510 + check_accuracy_budget([], [gap] * 30, tolerance, "systematically worse choices")
511 + check_accuracy_budget(
512 + [torch.tensor([0.03, 0.07])], [0.125, 0.125], tolerance, "control variation"
513 + )
514 +
515 +
516 +def _divergence_request(
517 + tokens: list[int], scores: dict[int, float], position: int = 0
518 +) -> RequestOutput:
519 + logprobs: list[dict[int, Logprob]] = [{} for _ in tokens]
520 + logprobs[position] = {
521 + token: Logprob(score, rank)
522 + for rank, (token, score) in enumerate(
523 + sorted(scores.items(), key=lambda item: item[1], reverse=True), 1
524 + )
525 + }
526 + return RequestOutput(
527 + request_id="test",
528 + prompt=None,
529 + prompt_token_ids=[0, 0, 0],
530 + prompt_logprobs=None,
531 + outputs=[
532 + CompletionOutput(
533 + index=0,
534 + text="",
535 + token_ids=tokens,
536 + cumulative_logprob=None,
537 + logprobs=logprobs,
538 + finish_reason="length",
539 + )
540 + ],
541 + finished=True,
542 + )
543 +
544 +
545 +@pytest.mark.parametrize(
546 + "base_scores,actual_scores,error",
547 + [
548 + ({0: -0.8, 1: -0.9}, {0: -0.9, 1: -0.8}, None),
549 + ({0: -0.8, 1: -0.8}, {0: -0.8, 1: -0.8}, None),
550 + ({0: -0.8, 1: -0.9}, {0: -0.8, 1: -0.9}, "actual.*below"),
551 + ({0: -0.9, 1: -0.8}, {0: -0.9, 1: -0.8}, "baseline.*below"),
552 + ({0: -0.8, 1: -1.1}, {0: -0.9, 1: -0.8}, "baseline.*near-tie"),
553 + ({0: -0.8, 1: -0.9}, {0: -1.1, 1: -0.8}, "actual.*near-tie"),
554 + ({0: -0.8, 2: -0.9}, {0: -0.9, 1: -0.8}, None),
555 + ({0: -0.8, 2: -1.1}, {0: -0.9, 1: -0.8}, "baseline.*lower bound"),
556 + ({0: -0.8, 1: -0.9}, {1: -0.8, 2: -1.1}, "actual.*lower bound"),
557 + ],
558 + ids=[
559 + "near-tie-rank-reversal",
560 + "exact-tie",
561 + "wrong-near-tie-actual-choice",
562 + "wrong-near-tie-baseline-choice",
563 + "decisive-baseline",
564 + "decisive-actual",
565 + "missing-candidate-inconclusive",
566 + "missing-candidate-decisive-baseline",
567 + "missing-candidate-decisive-actual",
568 + ],
569 +)
570 +def test_first_divergence_checks_engine_preferences(base_scores, actual_scores, error):
571 + baseline = _divergence_request([0], base_scores)
572 + actual = _divergence_request([1], actual_scores)
573 + args = (
574 + baseline,
575 + actual,
576 + SamplingParams(temperature=0, logprobs=2),
577 + AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25),
578 + "paired engines",
579 + )
580 + if error is None:
581 + check_first_divergence(*args)
582 + else:
583 + with pytest.raises(AssertionError, match=error):
584 + check_first_divergence(*args)
585 +
586 +
587 +@pytest.mark.parametrize("missing_alternative", [False, True])
588 +def test_first_divergence_uses_generated_penalty_history(missing_alternative):
589 + # Raw ordering favors token 1; one generated 0 reverses the penalized order.
590 + # The prompt contains three more 0s which must not contribute to penalties.
591 + base_scores = {0: -1.5, 1: -0.6}
592 + if missing_alternative:
593 + # The selected 0 is outside top-2. Use the top-k boundary (-0.6),
594 + # excluding the additional sampled token (-1.5), for the missing 1.
595 + base_scores = {0: -1.5, 2: -0.5, 3: -0.6}
596 + baseline = _divergence_request([0, 0, 2], base_scores, position=1)
597 + actual = _divergence_request([0, 1, 3, 4], {0: -1.6, 1: -0.5}, position=1)
598 + check_first_divergence(
599 + baseline,
600 + actual,
601 + SamplingParams(temperature=0, logprobs=2, frequency_penalty=-1),
602 + AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25),
603 + "different later histories are validated separately",
604 + )
605 +
606 +
607 +@pytest.mark.parametrize("mismatch", ["length", "finish_reason", "prompt"])
608 +def test_first_divergence_rejects_mismatches_without_a_token_change(mismatch):
609 + baseline = _divergence_request([0, 1], {})
610 + actual = _divergence_request([0, 1], {})
611 + if mismatch == "length":
612 + actual.outputs[0].token_ids = [0]
613 + elif mismatch == "finish_reason":
614 + actual.outputs[0].finish_reason = "stop"
615 + else:
616 + actual.prompt_token_ids = [1]
617 + with pytest.raises(AssertionError):
618 + check_first_divergence(
619 + baseline,
620 + actual,
621 + SamplingParams(temperature=0, logprobs=2),
622 + AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25),
623 + mismatch,
624 + )