| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from typing import Any |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm import SamplingParams |
| 9 | from vllm.config import CompilationConfig |
| 10 | |
| 11 | from ..utils import compute_acceptance_len |
| 12 | |
| 13 | |
| 14 | def load_and_process_dataset(data_name: str): |
| 15 | from datasets import load_dataset |
| 16 | |
| 17 | if data_name == "gsm8k": |
| 18 | dataset = load_dataset("openai/gsm8k", "main", split="test") |
| 19 | prompt_fmt = ( |
| 20 | "{question}\nPlease reason step by step," |
| 21 | " and put your final answer within \\boxed{{}}." |
| 22 | ) |
| 23 | dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) |
| 24 | elif data_name == "mt-bench": |
| 25 | dataset = load_dataset("HuggingFaceH4/mt_bench_prompts", split="train") |
| 26 | dataset = dataset.map(lambda x: {"turns": x["prompt"]}) |
| 27 | elif data_name == "humaneval": |
| 28 | dataset = load_dataset("openai/openai_humaneval", split="test") |
| 29 | prompt_fmt = ( |
| 30 | "Write a solution to the following problem and make sure" |
| 31 | " that it passes the tests:\n```python\n{prompt}\n```" |
| 32 | ) |
| 33 | dataset = dataset.map(lambda x: {"turns": [prompt_fmt.format(**x)]}) |
| 34 | |
| 35 | return dataset |
| 36 | |
| 37 | |
| 38 | def run_acceptance_length_eval( |
| 39 | monkeypatch: pytest.MonkeyPatch, |
| 40 | vllm_runner, |
| 41 | spec_config: dict[str, Any], |
| 42 | expected_acceptance_lengths: dict[str, float], |
| 43 | chat_template_kwargs: dict[str, Any], |
| 44 | use_mrv2: bool, |
| 45 | ): |
| 46 | """E2E acceptance-rate validation for speculative decoding. |
| 47 | |
| 48 | Drives one or more datasets (keyed in ``expected_acceptance_lengths``) |
| 49 | through the spec decode engine and asserts the mean acceptance length |
| 50 | stays within tolerance of the reference figure for each dataset. |
| 51 | """ |
| 52 | monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") |
| 53 | runner_config = spec_config.copy() |
| 54 | model = runner_config.pop("model") |
| 55 | runner_config.setdefault("block_size", None) |
| 56 | runner_config.setdefault("enable_chunked_prefill", None) |
| 57 | runner_config.setdefault("compilation_config", CompilationConfig()) |
| 58 | with vllm_runner(model, **runner_config) as spec_runner: |
| 59 | spec_llm = spec_runner.llm |
| 60 | # MT-Bench has 80 prompts and HumanEval has 164; truncate GSM8K. |
| 61 | max_prompts_per_dataset = 200 |
| 62 | |
| 63 | tokenizer = spec_llm.get_tokenizer() |
| 64 | for dataset_name, expected_len in expected_acceptance_lengths.items(): |
| 65 | dataset = load_and_process_dataset(dataset_name) |
| 66 | prev_metrics = None |
| 67 | acceptance_lengths = [] |
| 68 | for i in range(min(max_prompts_per_dataset, len(dataset))): |
| 69 | user_content = dataset[i]["turns"][0] |
| 70 | prompt_text = tokenizer.apply_chat_template( |
| 71 | [{"role": "user", "content": user_content}], |
| 72 | tokenize=False, |
| 73 | add_generation_prompt=True, |
| 74 | **chat_template_kwargs, |
| 75 | ) |
| 76 | |
| 77 | # Greedy (temp=0) so acceptance length is deterministic and comparable |
| 78 | # across runs. |
| 79 | spec_llm.generate( |
| 80 | [prompt_text], |
| 81 | SamplingParams(temperature=0, max_tokens=2048), |
| 82 | use_tqdm=False, |
| 83 | ) |
| 84 | current_metrics = spec_llm.get_metrics() |
| 85 | acceptance_len = compute_acceptance_len(current_metrics, prev_metrics) |
| 86 | prev_metrics = current_metrics |
| 87 | acceptance_lengths.append(acceptance_len) |
| 88 | |
| 89 | mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths) |
| 90 | # Fairly tight tolerance of 95% against the reference figures, |
| 91 | # watching for regressions. Can be relaxed if test is flaky but be sure to |
| 92 | # check for genuine issues such as #40727. |
| 93 | expected_len = expected_len * 0.95 |
| 94 | print( |
| 95 | f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" |
| 96 | f" (expected at least {expected_len:.2f})" |
| 97 | ) |
| 98 | |
| 99 | assert mean_acceptance_length >= expected_len, ( |
| 100 | f"acceptance_len for {dataset_name} is below expected threshold: " |
| 101 | f"{mean_acceptance_length:.2f} < {expected_len:.2f}" |
| 102 | ) |