A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 23:00 UTC
Test: added acceptance length tests (#32030)
Signed-off-by: rahul-tuli <rtuli@redhat.com>
Rahul Tuli committed
Jan 21, 2026 at 00:25 UTC
f0feb1cf81ef33573ffb603a4ff0d12e604a622f
2 files changed
+279
tests/v1/spec_decode/__init__.py
new
+2
@@ -0,0 +1,2 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
tests/v1/spec_decode/test_acceptance_length.py
new
+277
@@ -0,0 +1,277 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+EAGLE3 Acceptance Length Regression Tests.
5
+
6
+These tests verify that acceptance lengths for EAGLE3 speculative decoding
7
+do not regress across vLLM commits. Each test runs inference on the MT-Bench
8
+dataset and asserts that the mean acceptance length is within tolerance of
9
+the expected baseline.
10
+"""
11
+
12
+from dataclasses import dataclass, field
13
+from types import SimpleNamespace
14
+
15
+import pytest
16
+import torch
17
+
18
+from tests.conftest import VllmRunner
19
+from tests.utils import large_gpu_mark
20
+from vllm import SamplingParams
21
+from vllm.benchmarks.datasets import get_samples
22
+from vllm.inputs import TokensPrompt
23
+from vllm.platforms import current_platform
24
+from vllm.v1.attention.backends.registry import AttentionBackendEnum
25
+from vllm.v1.attention.selector import AttentionSelectorConfig
26
+from vllm.v1.metrics.reader import Counter, Vector
27
+
28
+
29
+@dataclass
30
+class Eagle3ModelConfig:
31
+ verifier: str
32
+ drafter: str
33
+ expected_acceptance_length: float
34
+ expected_acceptance_lengths_per_pos: list[float] = field(default_factory=list)
35
+ id: str = ""
36
+ # Backends that are incompatible with this model (will be skipped)
37
+ excluded_backends: set[AttentionBackendEnum] = field(default_factory=set)
38
+
39
+
40
+# Model configurations for EAGLE3 acceptance length tests.
41
+# Expected acceptance lengths are determined by running baseline benchmarks
42
+# using examples/offline_inference/spec_decode.py with the MT-Bench dataset.
43
+EAGLE3_MODEL_CONFIGS = [
44
+ Eagle3ModelConfig(
45
+ verifier="meta-llama/Llama-3.1-8B-Instruct",
46
+ drafter="RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3",
47
+ expected_acceptance_length=2.60,
48
+ expected_acceptance_lengths_per_pos=[0.7296, 0.5208, 0.3545],
49
+ id="llama3-8b-eagle3",
50
+ ),
51
+ Eagle3ModelConfig(
52
+ verifier="Qwen/Qwen3-8B",
53
+ drafter="RedHatAI/Qwen3-8B-speculator.eagle3",
54
+ expected_acceptance_length=2.26,
55
+ expected_acceptance_lengths_per_pos=[0.6541, 0.3993, 0.2020],
56
+ id="qwen3-8b-eagle3",
57
+ ),
58
+ Eagle3ModelConfig(
59
+ verifier="openai/gpt-oss-20b",
60
+ drafter="RedHatAI/gpt-oss-20b-speculator.eagle3",
61
+ expected_acceptance_length=2.56,
62
+ expected_acceptance_lengths_per_pos=[0.7165, 0.5120, 0.3337],
63
+ id="gpt-oss-20b-eagle3",
64
+ # FLASHINFER incompatible: gpt-oss-20b uses sink attention which
65
+ # FLASHINFER does not support ("sink setting not supported")
66
+ excluded_backends={AttentionBackendEnum.FLASHINFER},
67
+ ),
68
+]
69
+
70
+# Default test parameters
71
+DEFAULT_NUM_SPEC_TOKENS = 3
72
+DEFAULT_NUM_PROMPTS = 80
73
+DEFAULT_OUTPUT_LEN = 256
74
+DEFAULT_MAX_MODEL_LEN = 16384
75
+DEFAULT_RTOL = 0.05
76
+
77
+# TP sizes to test
78
+TP_SIZES = [1, 2, 4]
79
+
80
+
81
+# Backends excluded from testing due to significantly different behavior
82
+EXCLUDED_BACKENDS = {AttentionBackendEnum.FLEX_ATTENTION}
83
+
84
+
85
+def get_available_attention_backends() -> list[str]:
86
+ if not hasattr(current_platform, "get_valid_backends"):
87
+ return ["FLASH_ATTN"]
88
+
89
+ device_capability = current_platform.get_device_capability()
90
+ if device_capability is None:
91
+ return ["FLASH_ATTN"]
92
+
93
+ attn_selector_config = AttentionSelectorConfig(
94
+ head_size=128,
95
+ dtype=torch.bfloat16,
96
+ kv_cache_dtype=None,
97
+ block_size=None,
98
+ use_mla=False,
99
+ has_sink=False,
100
+ use_sparse=False,
101
+ use_mm_prefix=False,
102
+ )
103
+
104
+ valid_backends, _ = current_platform.get_valid_backends(
105
+ device_capability=device_capability,
106
+ attn_selector_config=attn_selector_config,
107
+ )
108
+
109
+ return [
110
+ backend.name
111
+ for backend, _ in valid_backends
112
+ if backend not in EXCLUDED_BACKENDS
113
+ ]
114
+
115
+
116
+def get_attention_backend_params() -> list[str]:
117
+ return get_available_attention_backends()
118
+
119
+
120
+def get_tp_size_params() -> list[pytest.param]:
121
+ num_gpus = torch.cuda.device_count() if torch.cuda.is_available() else 1
122
+ return [pytest.param(tp, id=f"tp{tp}") for tp in TP_SIZES if tp <= num_gpus]
123
+
124
+
125
+def get_mt_bench_prompts(
126
+ tokenizer, num_prompts: int = DEFAULT_NUM_PROMPTS
127
+) -> list[list[int]]:
128
+ args = SimpleNamespace(
129
+ dataset_name="hf",
130
+ dataset_path="philschmid/mt-bench",
131
+ num_prompts=num_prompts,
132
+ seed=42,
133
+ no_oversample=False,
134
+ endpoint_type="openai-chat",
135
+ input_len=None,
136
+ output_len=DEFAULT_OUTPUT_LEN,
137
+ sharegpt_output_len=DEFAULT_OUTPUT_LEN,
138
+ hf_name=None,
139
+ hf_split="train",
140
+ hf_subset=None,
141
+ hf_output_len=DEFAULT_OUTPUT_LEN,
142
+ no_stream=True,
143
+ disable_shuffle=False,
144
+ skip_chat_template=False,
145
+ )
146
+ samples = get_samples(args, tokenizer)
147
+ prompt_ids = [
148
+ tokenizer.encode(sample.prompt, add_special_tokens=False) for sample in samples
149
+ ]
150
+ return prompt_ids
151
+
152
+
153
+def extract_acceptance_metrics(metrics, num_spec_tokens: int) -> dict:
154
+ num_drafts = 0
155
+ num_accepted_tokens = 0
156
+ acceptance_counts = [0] * num_spec_tokens
157
+
158
+ for metric in metrics:
159
+ if metric.name == "vllm:spec_decode_num_drafts":
160
+ assert isinstance(metric, Counter)
161
+ num_drafts += metric.value
162
+ elif metric.name == "vllm:spec_decode_num_accepted_tokens":
163
+ assert isinstance(metric, Counter)
164
+ num_accepted_tokens += metric.value
165
+ elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
166
+ assert isinstance(metric, Vector)
167
+ for pos in range(min(len(metric.values), num_spec_tokens)):
168
+ acceptance_counts[pos] += metric.values[pos]
169
+
170
+ # Calculate mean acceptance length
171
+ # Formula: 1 + (accepted_tokens / num_drafts)
172
+ acceptance_length = 1 + (num_accepted_tokens / num_drafts) if num_drafts > 0 else 1
173
+
174
+ # Calculate per-position acceptance lengths (contribution to total)
175
+ # Each position contributes: accepted_at_pos / num_drafts
176
+ acceptance_lengths_per_pos = [
177
+ count / num_drafts if num_drafts > 0 else 0.0 for count in acceptance_counts
178
+ ]
179
+
180
+ return {
181
+ "acceptance_length": acceptance_length,
182
+ "acceptance_lengths_per_pos": acceptance_lengths_per_pos,
183
+ "num_drafts": num_drafts,
184
+ "num_accepted_tokens": num_accepted_tokens,
185
+ }
186
+
187
+
188
+@large_gpu_mark(min_gb=40)
189
+@pytest.mark.parametrize(
190
+ "model_config",
191
+ [pytest.param(config, id=config.id) for config in EAGLE3_MODEL_CONFIGS],
192
+)
193
+@pytest.mark.parametrize("num_spec_tokens", [DEFAULT_NUM_SPEC_TOKENS])
194
+@pytest.mark.parametrize("tp_size", get_tp_size_params())
195
+@pytest.mark.parametrize("attention_backend", get_attention_backend_params())
196
+def test_eagle3_acceptance_length(
197
+ model_config: Eagle3ModelConfig,
198
+ num_spec_tokens: int,
199
+ tp_size: int,
200
+ attention_backend: str,
201
+ monkeypatch: pytest.MonkeyPatch,
202
+):
203
+ # Skip if this backend is incompatible with the model
204
+ backend_enum = AttentionBackendEnum[attention_backend]
205
+ if backend_enum in model_config.excluded_backends:
206
+ pytest.skip(f"{attention_backend} is incompatible with {model_config.id}")
207
+
208
+ with monkeypatch.context() as m:
209
+ m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
210
+ m.setenv("VLLM_ATTENTION_BACKEND", attention_backend)
211
+
212
+ with VllmRunner(
213
+ model_name=model_config.verifier,
214
+ speculative_config={
215
+ "method": "eagle3",
216
+ "model": model_config.drafter,
217
+ "num_speculative_tokens": num_spec_tokens,
218
+ },
219
+ tensor_parallel_size=tp_size,
220
+ gpu_memory_utilization=0.7,
221
+ disable_log_stats=False,
222
+ max_model_len=DEFAULT_MAX_MODEL_LEN,
223
+ ) as vllm_runner:
224
+ tokenizer = vllm_runner.llm.get_tokenizer()
225
+ prompt_ids = get_mt_bench_prompts(tokenizer, DEFAULT_NUM_PROMPTS)
226
+
227
+ sampling_params = SamplingParams(
228
+ temperature=0,
229
+ max_tokens=DEFAULT_OUTPUT_LEN,
230
+ )
231
+ vllm_runner.llm.generate(
232
+ [TokensPrompt(prompt_token_ids=ids) for ids in prompt_ids],
233
+ sampling_params=sampling_params,
234
+ )
235
+
236
+ metrics = vllm_runner.llm.get_metrics()
237
+ results = extract_acceptance_metrics(metrics, num_spec_tokens)
238
+
239
+ actual_acceptance_length = results["acceptance_length"]
240
+ expected = model_config.expected_acceptance_length
241
+ actual_per_pos = results["acceptance_lengths_per_pos"]
242
+ expected_per_pos = model_config.expected_acceptance_lengths_per_pos
243
+
244
+ rel_error = abs(actual_acceptance_length - expected) / expected
245
+
246
+ assert rel_error <= DEFAULT_RTOL, (
247
+ f"Acceptance length regression detected for {model_config.id}!\n"
248
+ f" Expected: {expected:.3f}\n"
249
+ f" Actual: {actual_acceptance_length:.3f}\n"
250
+ f" Relative error: {rel_error:.2%} (tolerance: {DEFAULT_RTOL:.2%})\n"
251
+ f" Drafts: {results['num_drafts']}, "
252
+ f"Accepted tokens: {results['num_accepted_tokens']}"
253
+ )
254
+
255
+ if expected_per_pos and len(expected_per_pos) == len(actual_per_pos):
256
+ for pos, (actual, exp) in enumerate(
257
+ zip(actual_per_pos, expected_per_pos)
258
+ ):
259
+ if exp > 0:
260
+ pos_rel_error = abs(actual - exp) / exp
261
+ assert pos_rel_error <= DEFAULT_RTOL, (
262
+ f"Per-position acceptance length regression at pos {pos} "
263
+ f"for {model_config.id}!\n"
264
+ f" Expected: {exp:.3f}\n"
265
+ f" Actual: {actual:.3f}\n"
266
+ f" Relative error: {pos_rel_error:.2%} "
267
+ f"(tolerance: {DEFAULT_RTOL:.2%})"
268
+ )
269
+
270
+ print(
271
+ f"\n{model_config.id} [tp={tp_size}, backend={attention_backend}]: "
272
+ f"acceptance_length={actual_acceptance_length:.3f}"
273
+ f" (expected={expected:.3f}, rel_error={rel_error:.2%})"
274
+ )
275
+ print(f" Per-position: {[f'{v:.3f}' for v in actual_per_pos]}")
276
+ if expected_per_pos:
277
+ print(f" Expected: {[f'{v:.3f}' for v in expected_per_pos]}")