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 3, 20:15 UTC
Added peagle speculators support (#41826)
Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>
shanjiaz committed
May 12, 2026 at 10:55 UTC
6ccb10d7947697606fda5dc7b4a9fe2ca48fa2c7
6 files changed
+141
-52
.buildkite/test_areas/spec_decode.yaml
+5
-5
@@ -84,19 +84,19 @@ steps:
84
commands:
85
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
86
87
-- label: DFlash Speculators Correctness
88
- key: dflash-speculators-correctness
89
- timeout_in_minutes: 30
87
+- label: Speculators Correctness
88
+ key: speculators-correctness
89
+ timeout_in_minutes: 60
90
device: h100
91
optional: true
92
num_devices: 1
93
source_file_dependencies:
94
- vllm/v1/spec_decode/
95
- vllm/model_executor/models/qwen3_dflash.py
96
- - tests/v1/spec_decode/test_speculators_dflash.py
96
+ - tests/v1/spec_decode/test_speculators_correctness.py
97
commands:
98
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
99
- - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
99
+ - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
100
101
- label: Spec Decode MTP hybrid (B200)
102
timeout_in_minutes: 30
tests/models/registry.py
+15
@@ -1493,6 +1493,21 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1493
"Qwen/Qwen3-VL-8B-Instruct",
1494
speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
1495
),
1496
+ # [PEagle]
1497
+ "PEagleDraftModel": _HfExamplesInfo(
1498
+ "Qwen/Qwen3-8B",
1499
+ trust_remote_code=True,
1500
+ speculative_model="nm-testing/qwen3-8b-peagle-speculators",
1501
+ tokenizer="Qwen/Qwen3-8B",
1502
+ use_original_num_layers=True,
1503
+ ),
1504
+ "PeagleLlamaForCausalLM": _HfExamplesInfo(
1505
+ "Qwen/Qwen3-8B",
1506
+ trust_remote_code=True,
1507
+ speculative_model="nm-testing/qwen3-8b-peagle-speculators",
1508
+ tokenizer="Qwen/Qwen3-8B",
1509
+ use_original_num_layers=True,
1510
+ ),
1511
# [MTP]
1512
"DeepSeekMTPModel": _HfExamplesInfo(
1513
"luccafong/deepseek_mtp_main_random",
tests/v1/spec_decode/test_speculators_correctness.py
renamed
+84
-46
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import dataclasses
4
+
5
import pytest
6
import torch
7
@@ -9,22 +11,55 @@ from vllm import LLM
11
from vllm.config import SpeculativeConfig
12
from vllm.distributed import cleanup_dist_env_and_memory
13
12
-MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators"
13
-
14
-EXPECTED_GSM8K_ACCURACY = 0.885
15
-ACCURACY_RTOL = 0.03
16
-EXPECTED_ACCEPTANCE_LEN = 3.45
17
-ACCEPTANCE_LEN_RTOL = 0.15
18
-
19
-# Expected per-position acceptance rates (accepted_at_pos / num_drafts)
20
-# Based on GSM8K evaluation with Qwen3-8B dflash speculators.
21
-EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282]
22
-PER_POS_RTOL = 0.15
23
-
14
25
-def compute_spec_decode_stats(
26
- metrics,
27
-) -> dict:
15
+@dataclasses.dataclass
16
+class SpeculatorTestConfig:
17
+ model_path: str
18
+ method: str
19
+ display_name: str
20
+ expected_gsm8k_accuracy: float
21
+ accuracy_rtol: float
22
+ expected_acceptance_len: float
23
+ acceptance_len_rtol: float
24
+ expected_per_pos_acceptance_rates: tuple[float, ...]
25
+ per_pos_rtol: float
26
+ quantization: str | None = None
27
+ parallel_drafting: bool | None = None
28
+
29
+
30
+DFLASH_CONFIG = SpeculatorTestConfig(
31
+ model_path="nm-testing/dflash-qwen3-8b-speculators",
32
+ method="dflash",
33
+ display_name="DFlash",
34
+ expected_gsm8k_accuracy=0.885,
35
+ accuracy_rtol=0.03,
36
+ expected_acceptance_len=3.45,
37
+ acceptance_len_rtol=0.15,
38
+ expected_per_pos_acceptance_rates=(0.795, 0.611, 0.429, 0.282),
39
+ per_pos_rtol=0.15,
40
+ quantization="fp8",
41
+)
42
+
43
+PEAGLE_CONFIG = SpeculatorTestConfig(
44
+ model_path="nm-testing/qwen3-8b-peagle-speculators",
45
+ method="eagle3",
46
+ display_name="PEagle",
47
+ expected_gsm8k_accuracy=0.88,
48
+ accuracy_rtol=0.05,
49
+ expected_acceptance_len=2.27,
50
+ acceptance_len_rtol=0.20,
51
+ expected_per_pos_acceptance_rates=(0.66, 0.36, 0.18, 0.09),
52
+ per_pos_rtol=0.20,
53
+ parallel_drafting=True,
54
+)
55
+
56
+SPECULATOR_CONFIGS = [
57
+ pytest.param(DFLASH_CONFIG, id="dflash"),
58
+ pytest.param(PEAGLE_CONFIG, id="peagle"),
59
+]
60
+
61
+
62
+def compute_spec_decode_stats(metrics) -> dict:
63
"""Extract all spec-decode metrics and compute derived stats."""
64
name2metric = {m.name: m for m in metrics}
65
@@ -67,25 +102,26 @@ def print_spec_decode_stats(stats: dict) -> None:
102
print("===============================\n")
103
104
70
-def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
105
+@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
106
+def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config):
107
"""
72
- Test DFlash speculators model properly initializes speculative decoding.
108
+ Test speculators model properly initializes speculative decoding.
109
110
Verifies:
111
1. Speculative config is automatically initialized from speculators config
76
- 2. Method is detected as 'dflash'
77
- 3. The draft model path is correctly set
78
- 4. Speculative tokens count is valid (num_speculative_tokens=8)
79
- 5. Text generation works with speculative decoding enabled
112
+ 2. Method is detected correctly
113
+ 3. parallel_drafting is set correctly (if applicable)
114
+ 4. The draft model path is correctly set
115
+ 5. Speculative tokens count is valid
116
+ 6. Text generation works with speculative decoding enabled
117
"""
118
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
119
83
- with vllm_runner(
84
- MODEL_PATH,
85
- dtype=torch.bfloat16,
86
- enforce_eager=True,
87
- quantization="fp8",
88
- ) as vllm_model:
120
+ runner_kwargs = dict(dtype=torch.bfloat16, enforce_eager=True)
121
+ if config.quantization:
122
+ runner_kwargs["quantization"] = config.quantization
123
+
124
+ with vllm_runner(config.model_path, **runner_kwargs) as vllm_model:
125
vllm_config = vllm_model.llm.llm_engine.vllm_config
126
127
assert isinstance(vllm_config.speculative_config, SpeculativeConfig), (
@@ -93,40 +129,43 @@ def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
129
)
130
131
spec_config = vllm_config.speculative_config
96
- assert spec_config.method == "dflash", (
97
- f"Expected method='dflash', got '{spec_config.method}'"
132
+ assert spec_config.method == config.method, (
133
+ f"Expected method='{config.method}', got '{spec_config.method}'"
134
)
135
+ if config.parallel_drafting is not None:
136
+ assert spec_config.parallel_drafting is config.parallel_drafting, (
137
+ f"Expected parallel_drafting={config.parallel_drafting} "
138
+ f"for {config.display_name} model"
139
+ )
140
assert spec_config.num_speculative_tokens > 0, (
141
f"Expected positive speculative tokens, "
142
f"got {spec_config.num_speculative_tokens}"
143
)
103
- assert spec_config.model == MODEL_PATH, (
104
- f"Draft model should be {MODEL_PATH}, got {spec_config.model}"
144
+ assert spec_config.model == config.model_path, (
145
+ f"Draft model should be {config.model_path}, got {spec_config.model}"
146
)
147
148
vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20)
108
- assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}"
149
+ assert vllm_outputs, (
150
+ f"No outputs generated for speculators model {config.model_path}"
151
+ )
152
153
154
@pytest.mark.slow_test
155
@large_gpu_mark(min_gb=40)
113
-def test_dflash_speculators_correctness(monkeypatch):
156
+@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
157
+def test_speculators_correctness(monkeypatch, config):
158
"""
115
- E2E correctness test for DFlash via the speculators auto-detect path.
159
+ E2E correctness test via the speculators auto-detect path.
160
161
Evaluates GSM8k accuracy to ensure the speculators-format model produces
162
correct outputs, and checks that acceptance length does not collapse under
163
batched inference (lm-eval style).
120
-
121
- Observed per-position acceptance rates on GSM8K (1319 prompts):
122
- pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282,
123
- pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023
124
- Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128)
164
"""
165
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
166
167
spec_llm = LLM(
129
- model=MODEL_PATH,
168
+ model=config.model_path,
169
trust_remote_code=True,
170
max_model_len=4096,
171
max_num_seqs=128,
@@ -137,7 +176,7 @@ def test_dflash_speculators_correctness(monkeypatch):
176
177
results = evaluate_gsm8k_offline(spec_llm)
178
accuracy = results["accuracy"]
140
- accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL)
179
+ accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
180
assert accuracy >= accuracy_threshold, (
181
f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
182
)
@@ -147,19 +186,18 @@ def test_dflash_speculators_correctness(monkeypatch):
186
print_spec_decode_stats(stats)
187
188
acceptance_len = stats["acceptance_len"]
150
- al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL)
189
+ al_threshold = config.expected_acceptance_len * (1 - config.acceptance_len_rtol)
190
assert acceptance_len >= al_threshold, (
152
- f"DFlash speculators acceptance length too low: "
191
+ f"{config.display_name} speculators acceptance length too low: "
192
f"{acceptance_len:.2f} < {al_threshold:.2f}"
193
)
194
156
- # Check per-position acceptance rates for the first few positions.
195
per_pos_rates = stats["per_pos_acceptance_rates"]
158
- for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES):
196
+ for i, expected_rate in enumerate(config.expected_per_pos_acceptance_rates):
197
assert i < len(per_pos_rates), (
198
f"Missing per-position acceptance rate for position {i}"
199
)
162
- threshold = expected_rate * (1 - PER_POS_RTOL)
200
+ threshold = expected_rate * (1 - config.per_pos_rtol)
201
assert per_pos_rates[i] >= threshold, (
202
f"Per-position acceptance rate at pos {i} too low: "
203
f"{per_pos_rates[i]:.4f} < {threshold:.4f} "
vllm/model_executor/models/registry.py
+2
@@ -591,6 +591,8 @@ _SPECULATIVE_DECODING_MODELS = {
591
"EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
592
"EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
593
"DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
594
+ "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
595
+ "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
596
"Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
597
"Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
598
"LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
vllm/transformers_utils/configs/speculators/algos.py
+31
@@ -43,6 +43,37 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
43
]
44
45
46
+@register_speculator("peagle")
47
+def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
48
+ """
49
+ Apply PEagle (Parallel Eagle) specific configuration transformations to
50
+ the `dict` used to construct the Transformers PreTrainedConfig.
51
+
52
+ PEagle specific fields:
53
+ - draft_vocab_size: Size of the draft model's vocabulary
54
+ - target_hidden_size: Hidden size of the target model
55
+ - norm_before_residual: Whether to apply norm before residual connection
56
+ - norm_before_fc: Whether to apply RMSNorm before the fc projection
57
+ - mask_token_id (required): Token ID used for parallel drafting mask
58
+ placeholders, mapped to pard_token for the proposer
59
+ - eagle_aux_hidden_state_layer_ids: Layer indices from the target model
60
+ whose intermediate hidden states are used as auxiliary inputs
61
+ """
62
+ pre_trained_config["architectures"] = ["PeagleLlamaForCausalLM"]
63
+ pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
64
+ if config_dict.get("target_hidden_size") is not None:
65
+ pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
66
+ pre_trained_config["norm_before_residual"] = config_dict.get(
67
+ "norm_before_residual", False
68
+ )
69
+ pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
70
+ pre_trained_config["pard_token"] = config_dict["mask_token_id"]
71
+ if config_dict.get("eagle_aux_hidden_state_layer_ids"):
72
+ pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
73
+ "eagle_aux_hidden_state_layer_ids"
74
+ ]
75
+
76
+
77
@register_speculator("dflash")
78
def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
79
"""
vllm/transformers_utils/configs/speculators/base.py
+4
-1
@@ -131,7 +131,10 @@ class SpeculatorsConfig(PretrainedConfig):
131
)
132
133
# Build base vLLM speculative configuration
134
- return {
134
+ result = {
135
"method": config_dict.get("speculators_model_type"),
136
"num_speculative_tokens": num_speculative_tokens,
137
}
138
+ if result["method"] == "peagle":
139
+ result.update({"method": "eagle3", "parallel_drafting": True})
140
+ return result