[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>
This commit is contained in:
Andreas Karatzas
2026-10-01 01:06:24 +00:00
committed by GitHub
co-authored by OpenAI Codex OpenAI Codex
parent 37d6174096
commit df8fd42116
4 changed files with 1166 additions and 3 deletions
+29 -1
View File
@@ -454,7 +454,7 @@ steps:
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
commands:
- 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"
- 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"
- label: ":amd: (MI250) E2E Scheduling" # TBD
timeout_in_minutes: 55
@@ -470,6 +470,20 @@ steps:
commands:
- pytest -v -s v1/e2e/general/test_async_scheduling.py
- label: ":amd: (MI250) E2E Scheduling Accuracy" # TBD
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
dind: false
agent_pool: mi250_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/general/test_async_scheduling_accuracy.py
- label: ":amd: (MI250) V1 Engine" # TBD
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250]
@@ -2654,6 +2668,20 @@ steps:
commands:
- pytest -v -s v1/e2e/general/test_async_scheduling.py
- label: ":amd: (MI300) E2E Scheduling Accuracy" # TBD
timeout_in_minutes: 55
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
dind: false
agent_pool: mi300_1
optional: true
working_dir: "/vllm-workspace/tests"
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
- vllm/platforms/rocm.py
commands:
- pytest -v -s v1/e2e/general/test_async_scheduling_accuracy.py
- label: ":amd: (MI300) V1 Engine" # TBD
timeout_in_minutes: 45
mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300]
+22 -2
View File
@@ -85,6 +85,26 @@ steps:
source_file_dependencies:
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 18GB) E2E Scheduling Accuracy"
key: e2e-scheduling-accuracy-1-gpu
timeout_in_minutes: 53
device: h200_18gb
source_file_dependencies:
- vllm/v1/
- tests/v1/e2e/general/
commands:
- pytest -v -s v1/e2e/general/test_async_scheduling_accuracy.py
mirror:
amd:
label: ":amd: (MI250) E2E Scheduling Accuracy"
dind: false
device: mi250_1
timeout_in_minutes: 55
depends_on:
- image-build-amd
source_file_dependencies:
- vllm/platforms/rocm.py
- label: ":nvidia: (H200 MIG 18GB) E2E Core"
device: h200_18gb
key: e2e-core-1-gpu
@@ -93,13 +113,13 @@ steps:
- vllm/v1/
- tests/v1/e2e/general/
commands:
- 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"
- 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"
mirror:
amd:
label: ":amd: (MI250) E2E Core"
dind: false
commands:
- 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"
- 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"
device: mi250_1
timeout_in_minutes: 50
depends_on:
+491
View File
@@ -0,0 +1,491 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Numerical checks against an independent, teacher-forced target model."""
import json
import math
from collections import Counter
from dataclasses import dataclass
from functools import lru_cache
import torch
import xgrammar as xgr
from tokenizers.decoders import DecodeStream
from tests.conftest import HfRunner
from vllm import SamplingParams
from vllm.logprobs import Logprob
from vllm.outputs import RequestOutput
from vllm.tokenizers.detokenizer_utils import convert_ids_list_to_tokens
from vllm.v1.engine.logprobs import LogprobsProcessor
@dataclass(frozen=True)
class AccuracyTolerance:
logprob_atol: float
greedy_atol: float
logprob_mean_atol: float = 0.1
greedy_total_atol: float = 0.5
def logprob_error(self, reference: torch.Tensor) -> torch.Tensor:
return torch.full_like(reference, self.logprob_atol)
def check_first_divergence(
baseline: RequestOutput,
actual: RequestOutput,
params: SamplingParams,
tolerance: AccuracyTolerance,
context: str,
) -> None:
"""Require the two engines' scores to justify their first different choice."""
assert baseline.finished and actual.finished, context
assert len(baseline.outputs) == len(actual.outputs) == 1, context
assert baseline.prompt_token_ids is not None, context
assert baseline.prompt_token_ids == actual.prompt_token_ids, context
base, other = baseline.outputs[0], actual.outputs[0]
position = next(
(i for i, (a, b) in enumerate(zip(base.token_ids, other.token_ids)) if a != b),
None,
)
if position is None:
assert len(base.token_ids) == len(other.token_ids), (
f"{context}: one completion is a strict prefix of the other"
)
assert (base.finish_reason, base.stop_reason) == (
other.finish_reason,
other.stop_reason,
), f"{context}: identical tokens have different stopping reasons"
return
assert params.logprobs is not None and params.logprobs > 0, context
assert params.temperature == 0.0, context
assert params.presence_penalty == 0.0 and params.repetition_penalty == 1.0
assert not params.logit_bias, context
a, b = base.token_ids[position], other.token_ids[position]
counts = Counter(base.token_ids[:position])
context = (
f"{context}, first divergence position={position}, "
f"baseline token={a}, actual token={b}"
)
for side, completion, selected, alternative in (
("baseline", base, a, b),
("actual", other, b, a),
):
assert completion.logprobs is not None, context
assert len(completion.logprobs) == len(completion.token_ids), context
scores = completion.logprobs[position]
assert selected in scores, f"{context}, {side}: missing selected token score"
chosen = scores[selected].logprob
assert math.isfinite(chosen), f"{context}, {side}: non-finite selected score"
known = alternative in scores
if known:
competing = scores[alternative].logprob
assert math.isfinite(competing), (
f"{context}, {side}: divergent token has a non-finite score"
)
else:
# The sampled token may be returned in addition to the requested
# top-k. An omitted token's raw score is bounded by the kth score.
top_scores = [
score.logprob
for token, score in scores.items()
if len(scores) == params.logprobs or token != selected
]
assert len(top_scores) == params.logprobs, context
assert not any(math.isnan(score) for score in top_scores), context
competing = min(top_scores)
# Both candidates were emitted with the same history and constraints;
# frequency penalties are the only supported transform changing their
# finite score difference. Prompt tokens do not contribute to counts.
chosen -= params.frequency_penalty * counts[selected]
competing -= params.frequency_penalty * counts[alternative]
gap = chosen - competing
detail = (
f"{context}, {side}: selected={selected}, alternative={alternative}, "
f"{'gap' if known else 'gap lower bound'}={gap}"
)
if known:
# Log-softmax and reconstructing penalties may round in FP32, but
# a near tie cannot excuse selecting the engine's own worse token.
assert gap >= -1e-5, f"{detail}: selected token is below its alternative"
assert gap <= tolerance.greedy_atol, (
f"{detail}: divergence exceeds near-tie bound={tolerance.greedy_atol}"
)
def check_logprobs(
scores: dict[int, Logprob],
reference: torch.Tensor,
token: int,
requested: int,
tolerance: AccuracyTolerance,
context: str,
) -> torch.Tensor:
"""Check every score, reported rank, and the complete requested top-k."""
assert all(0 <= token_id < reference.numel() for token_id in scores), context
assert token in scores, f"{context}: missing score for token {token}"
assert max(1, requested) <= len(scores) <= requested + 1, context
assert not torch.isnan(reference).any(), context
assert all(s.rank is not None and s.rank >= 1 for s in scores.values()), context
ids = torch.tensor(list(scores), device=reference.device)
actual = torch.tensor(
[score.logprob for score in scores.values()], device=reference.device
)
expected = reference[ids]
finite = torch.isfinite(expected)
assert torch.equal(torch.isfinite(actual), finite), (
f"{context}: finite/masked logprobs differ: {scores}"
)
assert torch.equal(actual[~finite], expected[~finite]), context
assert (actual[finite] <= 0).all(), f"{context}: positive log probability"
errors = (actual[finite] - expected[finite]).abs()
bounds = tolerance.logprob_error(expected[finite])
assert (errors <= bounds).all(), (
f"{context}: token IDs={ids[finite].tolist()}, "
f"actual={actual[finite].tolist()}, reference={expected[finite].tolist()}, "
f"error={errors.tolist()}, bound={bounds.tolist()}"
)
for token_id, score in scores.items():
value = reference[token_id]
if not torch.isfinite(value):
continue
delta = tolerance.logprob_error(value)
lower = 1 + int((reference > value + 2 * delta).sum())
upper = max(1, int((reference >= value - 2 * delta).sum()))
assert score.rank is not None and lower <= score.rank <= upper, (
f"{context}: token={token_id}, rank={score.rank}, "
f"reference rank interval=[{lower}, {upper}]"
)
if requested:
# The sampled token can be outside top-k even when it ties the kth
# score. Its strict-greater rank can therefore duplicate a top-k rank.
top_ids = (
list(scores)
if len(scores) == requested
else [t for t in scores if t != token]
)
assert len(top_ids) == requested, f"{context}: incomplete top-k"
assert {scores[t].rank for t in top_ids} == set(range(1, requested + 1)), (
context
)
ordered = sorted(top_ids, key=lambda t: scores[t].rank)
assert all(
scores[a].logprob >= scores[b].logprob for a, b in zip(ordered, ordered[1:])
), f"{context}: top-k scores disagree with their reported ranks"
boundary = reference.topk(requested).values[-1]
if torch.isfinite(boundary):
delta = tolerance.logprob_error(boundary)
required = (reference > boundary + 2 * delta).nonzero().flatten()
assert set(required.tolist()) <= scores.keys(), (
f"{context}: missing an unambiguously top-{requested} token"
)
assert (reference[top_ids] >= boundary - 2 * delta).all(), (
f"{context}: returned token below the top-{requested} boundary"
)
return errors
def check_greedy_token(
logits: torch.Tensor,
token: int,
tolerance: AccuracyTolerance,
context: str,
) -> float:
assert not torch.isnan(logits).any(), context
assert torch.isfinite(logits[token]), (
f"{context}: selected token {token} is forbidden by the sampling parameters"
)
best = int(logits.argmax())
gap = float(logits[best] - logits[token])
assert gap <= tolerance.greedy_atol, (
f"{context}: selected token={token}, reference best={best}, "
f"logit gap={gap}, bound={tolerance.greedy_atol}"
)
return gap
def check_accuracy_budget(errors, gaps, tolerance, context):
"""Reject systematic sub-bound errors within each request and score stream."""
if errors:
finite_errors = torch.cat(errors)
if finite_errors.numel():
mean_error = float(finite_errors.mean())
assert mean_error <= tolerance.logprob_mean_atol, (
f"{context}: mean logprob error={mean_error}, "
f"bound={tolerance.logprob_mean_atol}"
)
assert sum(gaps) <= tolerance.greedy_total_atol, (
f"{context}: total greedy logit gap={sum(gaps)}, "
f"bound={tolerance.greedy_total_atol}"
)
def _check_decoded_tokens(processor, scores, history, context):
ids = list(scores)
decoded = convert_ids_list_to_tokens(processor.tokenizer, ids)
decoded = processor._verify_tokens(decoded, ids, context_token_ids=history[-4:])
for token, text in zip(ids, decoded, strict=True):
assert scores[token].decoded_token == text, (
f"{context}: incorrect decoded text for token {token}"
)
@torch.inference_mode()
def check_request_accuracy(
hf: HfRunner,
request: RequestOutput,
params: SamplingParams,
tolerance: AccuracyTolerance,
context: str,
) -> None:
"""Score the actual continuation, including positions after any divergence."""
assert request.finished and len(request.outputs) == 1, context
(completion,) = request.outputs
prompt = request.prompt_token_ids
assert prompt is not None and len(prompt) > 0, context
tokens = list(completion.token_ids)
assert params.max_tokens is not None and 0 < len(tokens) <= params.max_tokens, (
context
)
# Incremental decoding buffers incomplete UTF-8; decoding the final token
# list in one call can incorrectly append a replacement character.
stream = DecodeStream(skip_special_tokens=params.skip_special_tokens)
tokenizer = hf.tokenizer.backend_tokenizer
for token in prompt:
stream.step(tokenizer, token)
visible = tokens
if completion.finish_reason == "stop" and not params.include_stop_str_in_output:
visible = tokens[:-1]
expected_text = "".join(stream.step(tokenizer, token) or "" for token in visible)
assert completion.text == expected_text, (
f"{context}: decoded text does not match generated token IDs"
)
inputs = torch.tensor([prompt + tokens], device=hf.device)
logits = hf.model(input_ids=inputs, use_cache=False).logits[0].float()
# Use the existing byte-fallback handling with an independent request
# history; numerical scores always come from the HF model above.
decoder = LogprobsProcessor(
tokenizer=hf.tokenizer,
logprobs=None,
prompt_logprobs=None,
cumulative_logprob=None,
num_logprobs=None,
num_prompt_logprobs=None,
)
# Prompt token i is predicted by the causal logits at position i - 1.
if params.prompt_logprobs is None:
assert request.prompt_logprobs is None, context
else:
assert request.prompt_logprobs is not None, context
assert len(request.prompt_logprobs) == len(prompt), context
assert request.prompt_logprobs[0] is None, context
errors = []
for i, scores in enumerate(request.prompt_logprobs[1:], 1):
assert scores is not None, context
_check_decoded_tokens(decoder, scores, prompt[1:i], context)
errors.append(
check_logprobs(
scores,
logits[i - 1].log_softmax(-1),
prompt[i],
params.prompt_logprobs,
tolerance,
f"{context}, prompt position={i}",
)
)
check_accuracy_budget(errors, [], tolerance, f"{context}, prompt logprobs")
if params.logprobs is None:
assert completion.logprobs is None, context
else:
assert completion.logprobs is not None, context
assert len(completion.logprobs) == len(tokens), context
reference = SamplingReference(
hf.tokenizer,
logits.shape[-1],
hf.model.generation_config.to_dict(),
params,
)
errors = []
gaps = []
for i, token in enumerate(tokens):
raw, processed = reference.distributions(
logits[len(prompt) + i - 1], tokens[:i]
)
position = f"{context}, generated position={i}"
gaps.append(check_greedy_token(processed, token, tolerance, position))
if completion.logprobs is not None:
_check_decoded_tokens(decoder, completion.logprobs[i], tokens[:i], position)
errors.append(
check_logprobs(
completion.logprobs[i],
raw.log_softmax(-1),
token,
params.logprobs,
tolerance,
position,
)
)
check_accuracy_budget(errors, gaps, tolerance, f"{context}, completion")
# Advance the grammar through the last token too; an incomplete JSON prefix
# is allowed at max_tokens, but every emitted token must remain valid.
reference.accept_history(tokens)
check_stopping(
tokens,
completion.finish_reason,
reference.stop_ids,
params,
context,
eos_token_id=hf.tokenizer.eos_token_id,
stop_reason=completion.stop_reason,
)
def check_stopping(
tokens,
finish_reason,
stop_ids,
params,
context,
*,
eos_token_id: int | None,
stop_reason: int | str | None,
):
assert not params.ignore_eos and not params.stop, context
assert not stop_ids.intersection(tokens[:-1]), (
f"{context}: generation continued after a stop token"
)
if tokens[-1] in stop_ids:
assert finish_reason == "stop", (
f"{context}: stop token did not terminate output"
)
expected_stop_reason = None if tokens[-1] == eos_token_id else tokens[-1]
assert stop_reason == expected_stop_reason, (
f"{context}: stop reason={stop_reason}, expected={expected_stop_reason}"
)
else:
assert finish_reason == "length" and len(tokens) == params.max_tokens, context
assert stop_reason is None, f"{context}: length cap has a stop reason"
@lru_cache(maxsize=4)
def _compiler(tokenizer, vocab_size: int):
info = xgr.TokenizerInfo.from_huggingface(tokenizer, vocab_size=vocab_size)
return xgr.GrammarCompiler(info, max_threads=1, cache_enabled=True)
class SamplingReference:
"""One request's reference; histories must grow without changing their prefix."""
def __init__(
self,
tokenizer,
vocab_size: int,
generation_config: dict,
params: "SamplingParams",
):
assert params.temperature == 0.0, "Only the original greedy matrix is supported"
assert params.presence_penalty == 0.0 and params.repetition_penalty == 1.0
assert params.allowed_token_ids is None and not params.logit_bias
self.vocab_size = vocab_size
self.min_tokens = params.min_tokens
self.frequency_penalty = params.frequency_penalty
self.stop_ids = set(params.stop_token_ids or [])
for eos in (tokenizer.eos_token_id, generation_config.get("eos_token_id")):
if eos is not None:
self.stop_ids.update([eos] if isinstance(eos, int) else eos)
assert all(0 <= token < vocab_size for token in self.stop_ids)
# Match public bad-word tokenization semantics without using sampler code.
self.bad_words: list[list[int]] = []
for word in params.bad_words or []:
plain = tokenizer.encode(word.lstrip(), add_special_tokens=False)
spaced = tokenizer.encode(" " + word.lstrip(), add_special_tokens=False)
assert plain, f"Empty bad-word tokenization: {word!r}"
self.bad_words.append(plain)
if spaced and len(spaced) == len(plain) and spaced[0] != plain[0]:
self.bad_words.append(spaced)
self.history: list[int] = []
self.matcher = None
self.bitmask = None
if params.structured_outputs is not None:
schema = params.structured_outputs.json
assert schema is not None, "Only the original JSON schema is supported"
# Preserve dictionary order: xgrammar's default grammar orders fields.
schema = schema if isinstance(schema, str) else json.dumps(schema)
compiled = _compiler(tokenizer, vocab_size).compile_json_schema(
schema, any_whitespace=True
)
self.matcher = xgr.GrammarMatcher(
compiled, override_stop_tokens=sorted(self.stop_ids) or None
)
self.bitmask = xgr.allocate_token_bitmask(1, vocab_size)
def accept_history(self, generated_ids: list[int]) -> bool:
"""Validate newly generated grammar tokens, including the final EOS.
Return whether the grammar has terminated. Use this after the final
output token; no next-token distribution exists after termination.
"""
assert generated_ids[: len(self.history)] == self.history, (
"Each reference belongs to one monotonically growing request history"
)
if self.matcher is not None:
for token in generated_ids[len(self.history) :]:
assert self.matcher.accept_token(token), (
f"Grammar rejected token {token} at position {len(self.history)}"
)
self.history.append(token)
return self.matcher.is_terminated()
else:
self.history = generated_ids.copy()
return False
def distributions(
self, raw_logits: torch.Tensor, generated_ids: list[int]
) -> tuple[torch.Tensor, torch.Tensor]:
"""Return FP32 logits before and after non-grammar sampling transforms.
``raw_logits`` scores the next token after the prompt and
``generated_ids``. Returned tensors do not alias the model's tensor.
Normalize the first result to check reported completion logprobs;
the second supplies the greedy decision reference.
"""
assert raw_logits.shape == (self.vocab_size,)
assert not self.accept_history(generated_ids), (
"Cannot generate a distribution after grammar EOS"
)
grammar_logits = raw_logits.float().clone()
if self.matcher is not None:
assert self.bitmask is not None
self.matcher.fill_next_token_bitmask(self.bitmask)
words = self.bitmask[0, :, None]
allowed = ((words >> torch.arange(32)) & 1).flatten()
allowed = allowed[: self.vocab_size].bool().to(raw_logits.device)
grammar_logits.masked_fill_(~allowed, -torch.inf)
processed = grammar_logits.clone()
for bad_word in self.bad_words:
prefix = bad_word[:-1]
if not prefix or generated_ids[-len(prefix) :] == prefix:
processed[bad_word[-1]] = -torch.inf
if len(generated_ids) < self.min_tokens and self.stop_ids:
stop_ids = sorted(self.stop_ids)
stop_scores = processed[stop_ids].clone()
processed[stop_ids] = -torch.inf
# A completed grammar may have no continuation except a stop token.
if self.matcher is not None and torch.isneginf(processed).all():
processed[stop_ids] = stop_scores
if self.frequency_penalty:
for token, count in Counter(generated_ids).items():
processed[token] -= self.frequency_penalty * count
return grammar_logits, processed
@@ -0,0 +1,624 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Accuracy with automatic attention selection and the model's normal dtype.
Complements the strict Flex/FP32 invariance checks in test_async_scheduling.py.
"""
import os
from dataclasses import replace
from typing import Any
import pytest
import torch
from tests.conftest import HfRunner, VllmRunner
from tests.utils import (
large_gpu_mark,
single_gpu_only,
)
from tests.v1.e2e.general.accuracy_utils import (
AccuracyTolerance,
check_accuracy_budget,
check_first_divergence,
check_greedy_token,
check_logprobs,
check_request_accuracy,
check_stopping,
)
from vllm import SamplingParams
from vllm.logprobs import Logprob
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.platforms import current_platform
from vllm.sampling_params import StructuredOutputsParams
from vllm.utils.torch_utils import set_default_torch_num_threads
from vllm.v1.metrics.reader import Metric
MODEL = "Qwen/Qwen3-0.6B"
EAGLE3_TARGET_MODEL = "meta-llama/Llama-3.2-1B-Instruct"
# Need to enforce eager for MRV2 while we sort out cudagraph issues.
ENFORCE_EAGER = os.getenv("ENFORCE_EAGER", "0") == "1"
first_prompt = (
"The following numbers of the sequence "
+ ", ".join(str(i) for i in range(10))
+ " are:"
)
example_prompts = [first_prompt, "In one word, the capital of France is "] + [
f"Tell me about the number {i}: " for i in range(32)
]
default_params = dict(
temperature=0.0, # greedy
max_tokens=30,
min_tokens=28,
)
@single_gpu_only
def test_without_spec_decoding(
vllm_runner: type[VllmRunner],
hf_runner: type[HfRunner],
sample_json_schema,
):
"""Check target-model accuracy across scheduling and executor configurations."""
struct_outputs = StructuredOutputsParams(json=sample_json_schema)
test_sampling_params: list[dict[str, Any]] = [
dict(),
dict(frequency_penalty=-1.0),
dict(bad_words=["the", " the"]),
dict(bad_words=["the", " the"], logprobs=2),
dict(logprobs=2),
dict(logprobs=2, frequency_penalty=-1.0),
dict(prompt_logprobs=2),
dict(prompt_logprobs=2, logprobs=2),
dict(structured_outputs=struct_outputs),
dict(
structured_outputs=struct_outputs,
logprobs=2,
),
dict(
structured_outputs=struct_outputs,
frequency_penalty=-1.0,
),
dict(
structured_outputs=struct_outputs,
logprobs=2,
frequency_penalty=-1.0,
),
]
# test_preemption, executor, async_scheduling,
# spec_config, test_prefill_chunking
test_configs = [
(False, "mp", False, None, False),
(True, "mp", False, None, True),
(False, "mp", True, None, False),
(False, "uni", True, None, False),
(True, "mp", True, None, False),
(True, "uni", True, None, False),
(False, "mp", True, None, True),
(True, "mp", True, None, True),
(True, "uni", True, None, True),
]
run_tests(vllm_runner, hf_runner, MODEL, test_configs, test_sampling_params)
@single_gpu_only
@large_gpu_mark(min_gb=16)
def test_with_eagle3_spec_decoding(
vllm_runner: type[VllmRunner],
hf_runner: type[HfRunner],
sample_json_schema,
):
"""Check accuracy and acceptance, including drafts that exceed model length."""
spec_config = {
"method": "eagle3",
"num_speculative_tokens": 2,
"model": "nm-testing/Llama3_2_1B_speculator.eagle3",
}
# Set small draft model len to force doesn't-fit-in-drafter case.
spec_config_short = spec_config | {"max_model_len": 50}
struct_outputs = StructuredOutputsParams(json=sample_json_schema)
test_sampling_params: list[dict[str, Any]] = [
dict(),
dict(frequency_penalty=-1.0),
dict(bad_words=["the", " the"]),
dict(bad_words=["the", " the"], logprobs=2),
dict(logprobs=2),
dict(logprobs=2, frequency_penalty=-1.0),
dict(prompt_logprobs=2),
dict(prompt_logprobs=2, logprobs=2),
dict(structured_outputs=struct_outputs),
dict(
structured_outputs=struct_outputs,
logprobs=2,
frequency_penalty=-1.0,
),
]
# test_preemption, executor, async_scheduling,
# spec_config, test_prefill_chunking
test_configs = [
(False, "mp", False, None, False),
(False, "mp", False, spec_config, False),
(True, "mp", False, spec_config, True),
(True, "uni", False, spec_config_short, True),
(False, "mp", True, spec_config, False),
(True, "mp", True, spec_config, False),
(False, "mp", True, spec_config_short, True),
(True, "uni", True, spec_config, False),
(True, "uni", True, spec_config_short, False),
(True, "mp", True, spec_config, True),
(True, "uni", True, spec_config_short, True),
]
run_tests(
vllm_runner, hf_runner, EAGLE3_TARGET_MODEL, test_configs, test_sampling_params
)
@single_gpu_only
@pytest.mark.skipif(
current_platform.is_xpu(),
reason=("XPU matmul/attention kernels are not batch-invariant"),
)
def test_with_ngram_gpu_spec_decoding(
vllm_runner: type[VllmRunner], hf_runner: type[HfRunner]
):
"""Check ngram accuracy and acceptance across scheduling configurations."""
ngram_gpu_config = {
"method": "ngram_gpu",
"num_speculative_tokens": 3,
"prompt_lookup_max": 3,
"prompt_lookup_min": 2,
}
# test_preemption, executor, async_scheduling,
# spec_config, test_prefill_chunking
test_configs = [
(False, "mp", False, None, False),
(False, "mp", False, ngram_gpu_config, False),
(True, "mp", False, ngram_gpu_config, True),
(False, "mp", True, ngram_gpu_config, False),
(True, "mp", True, ngram_gpu_config, False),
(True, "uni", True, ngram_gpu_config, False),
(True, "mp", True, ngram_gpu_config, True),
]
run_tests(vllm_runner, hf_runner, MODEL, test_configs, [{}, dict(logprobs=2)])
def run_tests(
vllm_runner: type[VllmRunner],
hf_runner: type[HfRunner],
model: str,
test_configs: list[tuple],
test_sampling_params: list[dict[str, Any]],
):
"""Compare the shared prefix and independently validate every continuation."""
assert test_configs[0] == (False, "mp", False, None, False), (
"The baseline must use synchronous target-only decoding"
)
outputs = []
for n, config in enumerate(test_configs, 1):
outputs.append(
run_test(
vllm_runner,
model,
f"{n}/{len(test_configs)}",
test_sampling_params,
*config,
)
)
baseline_config, baseline_batches, _ = outputs[0]
# Compare vLLM's own decisions at the first divergence, where histories
# still match. HF then scores each continuation, including after divergence.
# Neither engine forces an attention backend, dtype, or matmul precision.
# Qwen/Llama BF16 controls and Qwen holdouts measured max score error <0.38,
# per-stream request mean <0.073, max greedy gap 0.125, and request sum 0.25.
# Per-request budgets also reject systematic errors below the scalar bounds.
tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
# These short CPU reductions cost more to distribute across a thread pool.
with set_default_torch_num_threads(1), hf_runner(model) as hf:
hf.model.eval()
for config, batches, _ in outputs:
assert len(batches) == len(test_sampling_params)
for baseline_batch, batch, overrides in zip(
baseline_batches, batches, test_sampling_params, strict=True
):
assert len(batch) == len(example_prompts), config
params = SamplingParams(**default_params, **overrides)
for i, (baseline, request, prompt) in enumerate(
zip(baseline_batch, batch, example_prompts, strict=True)
):
assert request.prompt_token_ids == hf.tokenizer.encode(prompt)
context = f"config=[{config}], params={overrides}, request={i}"
if params.logprobs is not None:
check_first_divergence(
baseline,
request,
params,
tolerance,
f"baseline=[{baseline_config}], {context}",
)
check_request_accuracy(
hf,
request,
params,
tolerance,
context,
)
print(f"ACCURACY PASSED: config=[{config}], params={overrides}")
baseline_acceptances = next((o[2] for o in outputs if o[2] is not None), None)
if baseline_acceptances is not None:
for settings, (config, _, acceptances) in zip(
test_configs, outputs, strict=True
):
if acceptances is None:
continue
test_preemption, _, _, spec_config, _ = settings
assert spec_config is not None
for baseline, actual, params in zip(
baseline_acceptances, acceptances, test_sampling_params, strict=True
):
context = f"config=[{config}], params={params}"
if spec_config.get("max_model_len") is None:
# Keep the original acceptance-quality floor per batch.
relative_drop = (
0.10 if current_platform.is_rocm() and test_preemption else 0.05
)
assert actual >= baseline * (1 - relative_drop), (
f"{context}: acceptance={actual}, baseline={baseline}"
)
else:
assert actual > 0.1, f"{context}: acceptance={actual}"
def run_test(
vllm_runner: type[VllmRunner],
model: str,
test_str: str,
sampling_param_tests: list[dict[str, Any]],
test_preemption: bool,
executor: str,
async_scheduling: bool,
spec_config: dict[str, Any] | None,
test_prefill_chunking: bool,
):
spec_decoding = spec_config is not None
spec_method = (spec_config or {}).get("method", "none")
# Chunked ngram decoding admits fewer concurrent requests under the
# 48-token budget. Its original 33-block cache did not preempt on ROCm.
# A 17-block cache forces contention while every original request fits
# comfortably within 256 tokens (the longest prompt plus output is 67).
cache_blocks = 17 if test_prefill_chunking and spec_method == "ngram_gpu" else 33
cache_arg: dict[str, Any] = (
dict(
num_gpu_blocks_override=cache_blocks,
max_model_len=(cache_blocks - 1) * 16,
)
if test_preemption
else dict(gpu_memory_utilization=0.9, max_model_len=4096)
)
spec_mml = (spec_config or {}).get("max_model_len")
test_config = (
f"executor={executor}, preemption={test_preemption}, "
f"async_sched={async_scheduling}, "
f"chunk_prefill={test_prefill_chunking}, "
f"spec_decoding={spec_decoding}, spec_method={spec_method}, spec_mml={spec_mml}"
)
print("-" * 80)
print(f"---- TESTING {test_str}: {test_config}")
print("-" * 80)
with vllm_runner(
model,
enable_chunked_prefill=test_prefill_chunking,
# Force prefill chunking
max_num_batched_tokens=48 if test_prefill_chunking else None,
enforce_eager=ENFORCE_EAGER,
async_scheduling=async_scheduling,
distributed_executor_backend=executor,
speculative_config=spec_config,
disable_log_stats=False,
enable_prefix_caching=False if current_platform.is_rocm() else None,
**cache_arg,
) as vllm_model:
results = []
acceptance_rates: list[float] | None = [] if spec_decoding else None
for override_params in sampling_param_tests:
metrics_before = vllm_model.llm.get_metrics()
print(f"----------- RUNNING PARAMS: {override_params}")
results.append(
vllm_model.llm.generate(
vllm_model.get_inputs(example_prompts),
sampling_params=SamplingParams(**default_params, **override_params),
)
)
metrics_after = vllm_model.llm.get_metrics()
if acceptance_rates is not None:
acceptance_rate = _get_acceptance_rate(metrics_before, metrics_after)
acceptance_rates.append(acceptance_rate)
print(f"ACCEPTANCE RATE {acceptance_rate}")
if test_preemption:
preemptions = _get_count(
metrics_before, metrics_after, "vllm:num_preemptions"
)
assert preemptions > 0, "preemption test had no preemptions"
# Preserve the original parameter-effect checks across repeated batches.
if len(results) > 1:
baseline = _result_signature(results[0])
for outputs, params in zip(results[1:], sampling_param_tests[1:], strict=True):
assert _result_signature(outputs) != baseline, (
f"{test_config}: sampling parameters had no observable effect: {params}"
)
# The old ROCm filter used schema-only as its first parameter set.
# Retain those comparisons after restoring the ordinary cases too.
schema_batches = [
batch
for batch, params in zip(results, sampling_param_tests, strict=True)
if params.get("structured_outputs") is not None
]
for batch in schema_batches[1:]:
assert _result_signature(batch) != _result_signature(schema_batches[0]), (
f"{test_config}: structured sampling parameters had no effect"
)
return test_config, results, acceptance_rates
def _result_signature(requests):
return [
(
list(request.outputs[0].token_ids),
request.outputs[0].text,
request.prompt_logprobs,
request.outputs[0].logprobs,
)
for request in requests
]
def _get_acceptance_rate(before: list[Metric], after: list[Metric]) -> float:
draft = _get_count(before, after, "vllm:spec_decode_num_draft_tokens")
accept = _get_count(before, after, "vllm:spec_decode_num_accepted_tokens")
assert draft > 0, "speculative test did not draft any tokens"
assert 0 <= accept <= draft, f"invalid acceptance counters: {accept=}, {draft=}"
return accept / draft
def _get_count(before: list[Metric], after: list[Metric], name: str) -> int:
before_val = next(m.value for m in before if m.name == name)
after_val = next(m.value for m in after if m.name == name)
return after_val - before_val
@pytest.mark.parametrize(
"corruption", ["score", "rank", "top_k", "missing", "nan", "positive", "order"]
)
def test_accuracy_checks_reject_incorrect_logprobs(corruption):
second = 2.875 if corruption == "order" else 1.0
reference = torch.tensor([3.0, second, -2.0, -5.0]).log_softmax(-1)
tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
scores = {
0: Logprob(logprob=float(reference[0]), rank=1, decoded_token="a"),
1: Logprob(logprob=float(reference[1]), rank=2, decoded_token="b"),
}
check_logprobs(scores, reference, 0, 2, tolerance, "control")
if corruption == "score":
scores[0] = replace(scores[0], logprob=scores[0].logprob - 1.0)
elif corruption == "rank":
scores[1] = replace(scores[1], rank=1)
elif corruption == "top_k":
del scores[1]
scores[3] = Logprob(logprob=float(reference[3]), rank=2, decoded_token="d")
elif corruption == "missing":
del scores[0]
elif corruption == "positive":
scores[0] = replace(scores[0], logprob=0.1)
elif corruption == "order":
scores[0] = replace(scores[0], rank=2)
scores[1] = replace(scores[1], rank=1)
else:
scores[0] = replace(scores[0], logprob=float("nan"))
with pytest.raises(AssertionError):
check_logprobs(scores, reference, 0, 2, tolerance, corruption)
def test_accuracy_checks_distinguish_ties_from_wrong_tokens():
tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
logits = torch.tensor([3.0, 2.875, -2.0, -torch.inf])
check_greedy_token(logits, 0, tolerance, "best")
check_greedy_token(logits, 1, tolerance, "near tie")
for token in (2, 3):
with pytest.raises(AssertionError):
check_greedy_token(logits, token, tolerance, "incorrect token")
# top-k may omit a sampled token tied at its boundary. The sampled rank
# can duplicate a sequential top-k rank without corrupting the result.
reference = torch.zeros(4).log_softmax(-1)
scores = {
token: Logprob(logprob=float(reference[token]), rank=rank)
for token, rank in [(3, 1), (0, 1), (1, 2)]
}
check_logprobs(scores, reference, 3, 2, tolerance, "tied top-k")
# A grammar with one allowed token still returns the requested top-k,
# including masked entries whose score is exactly negative infinity.
reference = torch.tensor([0.0, -torch.inf, -torch.inf])
scores = {0: Logprob(0.0, 1), 1: Logprob(-float("inf"), 2)}
check_logprobs(scores, reference, 0, 2, tolerance, "masked top-k")
def test_accuracy_checks_reject_incorrect_stopping():
params = SamplingParams(max_tokens=3)
stop_ids = {2, 3}
for tokens, reason, stop_reason in [
([0, 1, 2], "stop", None),
([0, 1, 3], "stop", 3),
([0, 1, 0], "length", None),
]:
check_stopping(
tokens,
reason,
stop_ids,
params,
"valid stopping at length cap",
eos_token_id=2,
stop_reason=stop_reason,
)
for tokens, reason, stop_reason in [
([0, 2, 1], "length", None),
([0, 1, 2], "length", None),
([0, 1, 2], "stop", 2),
([0, 1, 3], "stop", None),
([0, 1, 3], "stop", 2),
([0, 1, 0], "length", 3),
]:
with pytest.raises(AssertionError):
check_stopping(
tokens,
reason,
stop_ids,
params,
"incorrect stopping",
eos_token_id=2,
stop_reason=stop_reason,
)
def test_accuracy_checks_reject_systematic_sub_bound_errors():
tolerance = AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25)
reference = torch.tensor([3.0, 2.75, -2.0]).log_softmax(-1)
scores = {i: Logprob(float(reference[i]) - 0.4, i + 1) for i in (0, 1)}
# Each corrupted score fits the scalar bound, but the bias must not pass.
errors = check_logprobs(scores, reference, 0, 2, tolerance, "biased scores")
with pytest.raises(AssertionError, match="mean logprob error"):
check_accuracy_budget([errors], [], tolerance, "biased scores")
gap = check_greedy_token(reference, 1, tolerance, "one ambiguous choice")
with pytest.raises(AssertionError, match="total greedy logit gap"):
check_accuracy_budget([], [gap] * 30, tolerance, "systematically worse choices")
check_accuracy_budget(
[torch.tensor([0.03, 0.07])], [0.125, 0.125], tolerance, "control variation"
)
def _divergence_request(
tokens: list[int], scores: dict[int, float], position: int = 0
) -> RequestOutput:
logprobs: list[dict[int, Logprob]] = [{} for _ in tokens]
logprobs[position] = {
token: Logprob(score, rank)
for rank, (token, score) in enumerate(
sorted(scores.items(), key=lambda item: item[1], reverse=True), 1
)
}
return RequestOutput(
request_id="test",
prompt=None,
prompt_token_ids=[0, 0, 0],
prompt_logprobs=None,
outputs=[
CompletionOutput(
index=0,
text="",
token_ids=tokens,
cumulative_logprob=None,
logprobs=logprobs,
finish_reason="length",
)
],
finished=True,
)
@pytest.mark.parametrize(
"base_scores,actual_scores,error",
[
({0: -0.8, 1: -0.9}, {0: -0.9, 1: -0.8}, None),
({0: -0.8, 1: -0.8}, {0: -0.8, 1: -0.8}, None),
({0: -0.8, 1: -0.9}, {0: -0.8, 1: -0.9}, "actual.*below"),
({0: -0.9, 1: -0.8}, {0: -0.9, 1: -0.8}, "baseline.*below"),
({0: -0.8, 1: -1.1}, {0: -0.9, 1: -0.8}, "baseline.*near-tie"),
({0: -0.8, 1: -0.9}, {0: -1.1, 1: -0.8}, "actual.*near-tie"),
({0: -0.8, 2: -0.9}, {0: -0.9, 1: -0.8}, None),
({0: -0.8, 2: -1.1}, {0: -0.9, 1: -0.8}, "baseline.*lower bound"),
({0: -0.8, 1: -0.9}, {1: -0.8, 2: -1.1}, "actual.*lower bound"),
],
ids=[
"near-tie-rank-reversal",
"exact-tie",
"wrong-near-tie-actual-choice",
"wrong-near-tie-baseline-choice",
"decisive-baseline",
"decisive-actual",
"missing-candidate-inconclusive",
"missing-candidate-decisive-baseline",
"missing-candidate-decisive-actual",
],
)
def test_first_divergence_checks_engine_preferences(base_scores, actual_scores, error):
baseline = _divergence_request([0], base_scores)
actual = _divergence_request([1], actual_scores)
args = (
baseline,
actual,
SamplingParams(temperature=0, logprobs=2),
AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25),
"paired engines",
)
if error is None:
check_first_divergence(*args)
else:
with pytest.raises(AssertionError, match=error):
check_first_divergence(*args)
@pytest.mark.parametrize("missing_alternative", [False, True])
def test_first_divergence_uses_generated_penalty_history(missing_alternative):
# Raw ordering favors token 1; one generated 0 reverses the penalized order.
# The prompt contains three more 0s which must not contribute to penalties.
base_scores = {0: -1.5, 1: -0.6}
if missing_alternative:
# The selected 0 is outside top-2. Use the top-k boundary (-0.6),
# excluding the additional sampled token (-1.5), for the missing 1.
base_scores = {0: -1.5, 2: -0.5, 3: -0.6}
baseline = _divergence_request([0, 0, 2], base_scores, position=1)
actual = _divergence_request([0, 1, 3, 4], {0: -1.6, 1: -0.5}, position=1)
check_first_divergence(
baseline,
actual,
SamplingParams(temperature=0, logprobs=2, frequency_penalty=-1),
AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25),
"different later histories are validated separately",
)
@pytest.mark.parametrize("mismatch", ["length", "finish_reason", "prompt"])
def test_first_divergence_rejects_mismatches_without_a_token_change(mismatch):
baseline = _divergence_request([0, 1], {})
actual = _divergence_request([0, 1], {})
if mismatch == "length":
actual.outputs[0].token_ids = [0]
elif mismatch == "finish_reason":
actual.outputs[0].finish_reason = "stop"
else:
actual.prompt_token_ids = [1]
with pytest.raises(AssertionError):
check_first_divergence(
baseline,
actual,
SamplingParams(temperature=0, logprobs=2),
AccuracyTolerance(logprob_atol=0.5, greedy_atol=0.25),
mismatch,
)