| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """EAGLE3 Acceptance Length Regression Tests. |
| 4 | |
| 5 | These tests verify that acceptance lengths for EAGLE3 speculative decoding |
| 6 | do not regress across vLLM commits. Each test runs inference on the MT-Bench |
| 7 | dataset and asserts that the mean acceptance length is within tolerance of |
| 8 | the expected baseline. |
| 9 | """ |
| 10 | |
| 11 | from dataclasses import dataclass, field |
| 12 | from types import SimpleNamespace |
| 13 | |
| 14 | import pytest |
| 15 | import torch |
| 16 | |
| 17 | from tests.conftest import VllmRunner |
| 18 | from tests.utils import large_gpu_mark |
| 19 | from vllm import SamplingParams |
| 20 | from vllm.benchmarks.datasets import get_samples |
| 21 | from vllm.inputs import TokensPrompt |
| 22 | from vllm.platforms import current_platform |
| 23 | from vllm.v1.attention.backends.registry import AttentionBackendEnum |
| 24 | from vllm.v1.attention.selector import AttentionSelectorConfig |
| 25 | from vllm.v1.metrics.reader import Counter, Vector |
| 26 | |
| 27 | |
| 28 | @dataclass |
| 29 | class Eagle3ModelConfig: |
| 30 | verifier: str |
| 31 | drafter: str |
| 32 | expected_acceptance_length: float |
| 33 | expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list) |
| 34 | id: str = "" |
| 35 | # Backends that are incompatible with this model (will be skipped) |
| 36 | excluded_backends: set[AttentionBackendEnum] = field(default_factory=set) |
| 37 | # Pytest marks for this configuration |
| 38 | marks: list = field(default_factory=list) |
| 39 | # Custom relative tolerance (defaults to DEFAULT_RTOL if None) |
| 40 | rtol: float | None = None |
| 41 | # ROCm-specific test configuration |
| 42 | rocm_expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list) |
| 43 | |
| 44 | |
| 45 | # Model configurations for EAGLE3 acceptance length tests. |
| 46 | # Expected acceptance lengths are determined by running baseline benchmarks |
| 47 | # using examples/features/speculative_decoding/spec_decode_offline.py |
| 48 | # with the MT-Bench dataset. |
| 49 | EAGLE3_MODEL_CONFIGS = [ |
| 50 | Eagle3ModelConfig( |
| 51 | verifier="meta-llama/Llama-3.1-8B-Instruct", |
| 52 | drafter="RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", |
| 53 | expected_acceptance_length=2.60, |
| 54 | expected_acceptance_lengths_per_pos=[0.7296, 0.5208, 0.3545], |
| 55 | id="llama3-8b-eagle3", |
| 56 | ), |
| 57 | Eagle3ModelConfig( |
| 58 | verifier="Qwen/Qwen3-8B", |
| 59 | drafter="RedHatAI/Qwen3-8B-speculator.eagle3", |
| 60 | expected_acceptance_length=2.26, |
| 61 | expected_acceptance_lengths_per_pos=[0.6541, 0.3993, 0.2020], |
| 62 | id="qwen3-8b-eagle3", |
| 63 | ), |
| 64 | Eagle3ModelConfig( |
| 65 | verifier="openai/gpt-oss-20b", |
| 66 | drafter="RedHatAI/gpt-oss-20b-speculator.eagle3", |
| 67 | expected_acceptance_length=2.56, |
| 68 | expected_acceptance_lengths_per_pos=[0.7165, 0.5120, 0.3337], |
| 69 | id="gpt-oss-20b-eagle3", |
| 70 | # FLASHINFER incompatible: gpt-oss-20b uses sink attention which |
| 71 | # FLASHINFER does not support ("sink setting not supported") |
| 72 | excluded_backends={AttentionBackendEnum.FLASHINFER}, |
| 73 | rocm_expected_acceptance_lengths_per_pos=[0.7040, 0.4820, 0.3350], |
| 74 | ), |
| 75 | Eagle3ModelConfig( |
| 76 | verifier="Qwen/Qwen3-VL-30B-A3B-Instruct-FP8", |
| 77 | drafter="nm-testing/Speculator-Qwen3-30B-MOE-VL-Eagle3", |
| 78 | expected_acceptance_length=1.35, |
| 79 | expected_acceptance_lengths_per_pos=[0.2900, 0.0620, 0.0115], |
| 80 | id="qwen3-30b-moe-vl-eagle3", |
| 81 | marks=[ |
| 82 | pytest.mark.slow_test, |
| 83 | ], |
| 84 | rtol=0.15, # Higher tolerance due to small absolute values at position 2 |
| 85 | ), |
| 86 | ] |
| 87 | |
| 88 | # Default test parameters |
| 89 | DEFAULT_NUM_SPEC_TOKENS = 3 |
| 90 | DEFAULT_NUM_PROMPTS = 80 |
| 91 | DEFAULT_OUTPUT_LEN = 256 |
| 92 | DEFAULT_MAX_MODEL_LEN = 16384 |
| 93 | DEFAULT_RTOL = 0.05 |
| 94 | |
| 95 | # TP sizes to test |
| 96 | TP_SIZES = [1, 2, 4] |
| 97 | |
| 98 | |
| 99 | # Backends excluded from testing due to significantly different behavior |
| 100 | EXCLUDED_BACKENDS = {AttentionBackendEnum.FLEX_ATTENTION} |
| 101 | |
| 102 | |
| 103 | def get_available_attention_backends() -> list[str]: |
| 104 | if current_platform.is_rocm(): |
| 105 | return ["auto"] |
| 106 | |
| 107 | # Check if get_valid_backends is actually defined in the platform class |
| 108 | # (not just returning None from __getattr__) |
| 109 | get_valid_backends = getattr(current_platform.__class__, "get_valid_backends", None) |
| 110 | if get_valid_backends is None: |
| 111 | return ["FLASH_ATTN"] |
| 112 | |
| 113 | device_capability = current_platform.get_device_capability() |
| 114 | if device_capability is None: |
| 115 | return ["FLASH_ATTN"] |
| 116 | |
| 117 | attn_selector_config = AttentionSelectorConfig( |
| 118 | head_size=128, |
| 119 | dtype=torch.bfloat16, |
| 120 | kv_cache_dtype=None, |
| 121 | block_size=None, |
| 122 | use_mla=False, |
| 123 | has_sink=False, |
| 124 | use_sparse=False, |
| 125 | use_mm_prefix=False, |
| 126 | ) |
| 127 | |
| 128 | valid_backends, _ = current_platform.get_valid_backends( |
| 129 | device_capability=device_capability, |
| 130 | attn_selector_config=attn_selector_config, |
| 131 | ) |
| 132 | |
| 133 | return [ |
| 134 | candidate.backend.name |
| 135 | for candidate in valid_backends |
| 136 | if candidate.backend not in EXCLUDED_BACKENDS |
| 137 | ] |
| 138 | |
| 139 | |
| 140 | def get_attention_backend_params() -> list[str]: |
| 141 | return get_available_attention_backends() |
| 142 | |
| 143 | |
| 144 | def get_tp_size_params() -> list[pytest.param]: |
| 145 | num_gpus = torch.accelerator.device_count() if torch.cuda.is_available() else 1 |
| 146 | return [pytest.param(tp, id=f"tp{tp}") for tp in TP_SIZES if tp <= num_gpus] |
| 147 | |
| 148 | |
| 149 | def get_mt_bench_prompts( |
| 150 | tokenizer, num_prompts: int = DEFAULT_NUM_PROMPTS |
| 151 | ) -> list[list[int]]: |
| 152 | args = SimpleNamespace( |
| 153 | dataset_name="hf", |
| 154 | dataset_path="philschmid/mt-bench", |
| 155 | num_prompts=num_prompts, |
| 156 | seed=42, |
| 157 | no_oversample=False, |
| 158 | endpoint_type="openai-chat", |
| 159 | input_len=None, |
| 160 | output_len=DEFAULT_OUTPUT_LEN, |
| 161 | sharegpt_output_len=DEFAULT_OUTPUT_LEN, |
| 162 | hf_name=None, |
| 163 | hf_split="train", |
| 164 | hf_subset=None, |
| 165 | hf_output_len=DEFAULT_OUTPUT_LEN, |
| 166 | no_stream=True, |
| 167 | disable_shuffle=False, |
| 168 | skip_chat_template=False, |
| 169 | trust_remote_code=False, |
| 170 | enable_multimodal_chat=False, |
| 171 | request_id_prefix="", |
| 172 | ) |
| 173 | samples = get_samples(args, tokenizer) |
| 174 | prompt_ids = [ |
| 175 | tokenizer.encode(sample.prompt, add_special_tokens=False) for sample in samples |
| 176 | ] |
| 177 | return prompt_ids |
| 178 | |
| 179 | |
| 180 | def extract_acceptance_metrics(metrics, num_spec_tokens: int) -> dict: |
| 181 | num_drafts = 0 |
| 182 | num_accepted_tokens = 0 |
| 183 | acceptance_counts = [0] * num_spec_tokens |
| 184 | |
| 185 | for metric in metrics: |
| 186 | if metric.name == "vllm:spec_decode_num_drafts": |
| 187 | assert isinstance(metric, Counter) |
| 188 | num_drafts += metric.value |
| 189 | elif metric.name == "vllm:spec_decode_num_accepted_tokens": |
| 190 | assert isinstance(metric, Counter) |
| 191 | num_accepted_tokens += metric.value |
| 192 | elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos": |
| 193 | assert isinstance(metric, Vector) |
| 194 | for pos in range(min(len(metric.values), num_spec_tokens)): |
| 195 | acceptance_counts[pos] += metric.values[pos] |
| 196 | |
| 197 | # Calculate mean acceptance length |
| 198 | # Formula: 1 + (accepted_tokens / num_drafts) |
| 199 | acceptance_length = 1 + (num_accepted_tokens / num_drafts) if num_drafts > 0 else 1 |
| 200 | |
| 201 | # Calculate per-position acceptance lengths (contribution to total) |
| 202 | # Each position contributes: accepted_at_pos / num_drafts |
| 203 | acceptance_lengths_per_pos = [ |
| 204 | count / num_drafts if num_drafts > 0 else 0.0 for count in acceptance_counts |
| 205 | ] |
| 206 | |
| 207 | return { |
| 208 | "acceptance_length": acceptance_length, |
| 209 | "acceptance_lengths_per_pos": acceptance_lengths_per_pos, |
| 210 | "num_drafts": num_drafts, |
| 211 | "num_accepted_tokens": num_accepted_tokens, |
| 212 | } |
| 213 | |
| 214 | |
| 215 | @large_gpu_mark(min_gb=40) |
| 216 | @pytest.mark.skipif( |
| 217 | not current_platform.is_cuda_alike(), |
| 218 | reason="This test is only supported on CUDA-alike platforms.", |
| 219 | ) |
| 220 | @pytest.mark.parametrize( |
| 221 | "model_config", |
| 222 | [ |
| 223 | pytest.param(config, id=config.id, marks=config.marks) |
| 224 | for config in EAGLE3_MODEL_CONFIGS |
| 225 | ], |
| 226 | ) |
| 227 | @pytest.mark.parametrize("num_spec_tokens", [DEFAULT_NUM_SPEC_TOKENS]) |
| 228 | @pytest.mark.parametrize("tp_size", get_tp_size_params()) |
| 229 | @pytest.mark.parametrize("attention_backend", get_attention_backend_params()) |
| 230 | def test_eagle3_acceptance_length( |
| 231 | model_config: Eagle3ModelConfig, |
| 232 | num_spec_tokens: int, |
| 233 | tp_size: int, |
| 234 | attention_backend: str, |
| 235 | monkeypatch: pytest.MonkeyPatch, |
| 236 | ): |
| 237 | # Skip if this backend is incompatible with the model |
| 238 | attention_config = None |
| 239 | if attention_backend != "auto": |
| 240 | backend_enum = AttentionBackendEnum[attention_backend] |
| 241 | if backend_enum in model_config.excluded_backends: |
| 242 | pytest.skip(f"{attention_backend} is incompatible with {model_config.id}") |
| 243 | attention_config = {"backend": attention_backend} |
| 244 | |
| 245 | with monkeypatch.context() as m: |
| 246 | m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 247 | |
| 248 | with VllmRunner( |
| 249 | model_name=model_config.verifier, |
| 250 | speculative_config={ |
| 251 | "method": "eagle3", |
| 252 | "model": model_config.drafter, |
| 253 | "num_speculative_tokens": num_spec_tokens, |
| 254 | }, |
| 255 | attention_config=attention_config, |
| 256 | tensor_parallel_size=tp_size, |
| 257 | gpu_memory_utilization=0.7, |
| 258 | disable_log_stats=False, |
| 259 | max_model_len=DEFAULT_MAX_MODEL_LEN, |
| 260 | # Qwen/Qwen3-30B-A3B-FP8 with TP=4 needs EP |
| 261 | # https://github.com/vllm-project/vllm/issues/25292 |
| 262 | enable_expert_parallel=( |
| 263 | tp_size == 4 and "Qwen3-VL" in model_config.verifier |
| 264 | ), |
| 265 | ) as vllm_runner: |
| 266 | tokenizer = vllm_runner.llm.get_tokenizer() |
| 267 | prompt_ids = get_mt_bench_prompts(tokenizer, DEFAULT_NUM_PROMPTS) |
| 268 | |
| 269 | sampling_params = SamplingParams( |
| 270 | temperature=0, |
| 271 | max_tokens=DEFAULT_OUTPUT_LEN, |
| 272 | ) |
| 273 | vllm_runner.llm.generate( |
| 274 | [TokensPrompt(prompt_token_ids=ids) for ids in prompt_ids], |
| 275 | sampling_params=sampling_params, |
| 276 | ) |
| 277 | |
| 278 | metrics = vllm_runner.llm.get_metrics() |
| 279 | results = extract_acceptance_metrics(metrics, num_spec_tokens) |
| 280 | |
| 281 | actual_acceptance_length = results["acceptance_length"] |
| 282 | expected = model_config.expected_acceptance_length |
| 283 | actual_per_pos = results["acceptance_lengths_per_pos"] |
| 284 | expected_per_pos = model_config.expected_acceptance_lengths_per_pos |
| 285 | if ( |
| 286 | current_platform.is_rocm() |
| 287 | and model_config.rocm_expected_acceptance_lengths_per_pos |
| 288 | ): |
| 289 | expected_per_pos = model_config.rocm_expected_acceptance_lengths_per_pos |
| 290 | |
| 291 | rel_error = abs(actual_acceptance_length - expected) / expected |
| 292 | |
| 293 | # Overall acceptance length always uses DEFAULT_RTOL |
| 294 | assert rel_error <= DEFAULT_RTOL, ( |
| 295 | f"Acceptance length regression detected for {model_config.id}!\n" |
| 296 | f" Expected: {expected:.3f}\n" |
| 297 | f" Actual: {actual_acceptance_length:.3f}\n" |
| 298 | f" Relative error: {rel_error:.2%} (tolerance: {DEFAULT_RTOL:.2%})\n" |
| 299 | f" Drafts: {results['num_drafts']}, " |
| 300 | f"Accepted tokens: {results['num_accepted_tokens']}" |
| 301 | ) |
| 302 | |
| 303 | if expected_per_pos and len(expected_per_pos) == len(actual_per_pos): |
| 304 | # Per-position checks use model-specific rtol if provided |
| 305 | rtol = ( |
| 306 | model_config.rtol if model_config.rtol is not None else DEFAULT_RTOL |
| 307 | ) |
| 308 | for pos, (actual, exp) in enumerate( |
| 309 | zip(actual_per_pos, expected_per_pos) |
| 310 | ): |
| 311 | if exp > 0: |
| 312 | min_expected = exp * (1 - rtol) |
| 313 | assert actual >= min_expected, ( |
| 314 | f"Per-position acceptance length regression at pos {pos} " |
| 315 | f"for {model_config.id}!\n" |
| 316 | f" Expected: {exp:.3f}\n" |
| 317 | f" Actual: {actual:.3f}\n" |
| 318 | f" Minimum: {min_expected:.3f}\n" |
| 319 | f" Tolerance: rtol={rtol:.2%}" |
| 320 | ) |
| 321 | |
| 322 | print( |
| 323 | f"\n{model_config.id} [tp={tp_size}, backend={attention_backend}]: " |
| 324 | f"acceptance_length={actual_acceptance_length:.3f}" |
| 325 | f" (expected={expected:.3f}, rel_error={rel_error:.2%})" |
| 326 | ) |
| 327 | print(f" Per-position: {[f'{v:.3f}' for v in actual_per_pos]}") |
| 328 | if expected_per_pos: |
| 329 | print(f" Expected: {[f'{v:.3f}' for v in expected_per_pos]}") |