Added peagle speculators support (#41826)

Signed-off-by: shanjiaz <zsjwpianpian@gmail.com>

.buildkite/test_areas/spec_decode.yaml

@@ -84,19 +84,19 @@ steps:
   commands:
     - pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
 
-- label: DFlash Speculators Correctness
-  key: dflash-speculators-correctness
-  timeout_in_minutes: 30
+- label: Speculators Correctness
+  key: speculators-correctness
+  timeout_in_minutes: 60
   device: h100
   optional: true
   num_devices: 1
   source_file_dependencies:
     - vllm/v1/spec_decode/
     - vllm/model_executor/models/qwen3_dflash.py
-    - tests/v1/spec_decode/test_speculators_dflash.py
+    - tests/v1/spec_decode/test_speculators_correctness.py
   commands:
     - export VLLM_ALLOW_INSECURE_SERIALIZATION=1
-    - pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
+    - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test
 
 - label: Spec Decode MTP hybrid (B200)
   timeout_in_minutes: 30

tests/models/registry.py

@@ -1493,6 +1493,21 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
         "Qwen/Qwen3-VL-8B-Instruct",
         speculative_model="taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3",
     ),
+    # [PEagle]
+    "PEagleDraftModel": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        trust_remote_code=True,
+        speculative_model="nm-testing/qwen3-8b-peagle-speculators",
+        tokenizer="Qwen/Qwen3-8B",
+        use_original_num_layers=True,
+    ),
+    "PeagleLlamaForCausalLM": _HfExamplesInfo(
+        "Qwen/Qwen3-8B",
+        trust_remote_code=True,
+        speculative_model="nm-testing/qwen3-8b-peagle-speculators",
+        tokenizer="Qwen/Qwen3-8B",
+        use_original_num_layers=True,
+    ),
     # [MTP]
     "DeepSeekMTPModel": _HfExamplesInfo(
         "luccafong/deepseek_mtp_main_random",

tests/v1/spec_decode/test_speculators_correctness.py

@@ -1,5 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import dataclasses
+
 import pytest
 import torch
 
@@ -9,22 +11,55 @@ from vllm import LLM
 from vllm.config import SpeculativeConfig
 from vllm.distributed import cleanup_dist_env_and_memory
 
-MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators"
 
-EXPECTED_GSM8K_ACCURACY = 0.885
-ACCURACY_RTOL = 0.03
-EXPECTED_ACCEPTANCE_LEN = 3.45
-ACCEPTANCE_LEN_RTOL = 0.15
-
-# Expected per-position acceptance rates (accepted_at_pos / num_drafts)
-# Based on GSM8K evaluation with Qwen3-8B dflash speculators.
-EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282]
-PER_POS_RTOL = 0.15
+@dataclasses.dataclass
+class SpeculatorTestConfig:
+    model_path: str
+    method: str
+    display_name: str
+    expected_gsm8k_accuracy: float
+    accuracy_rtol: float
+    expected_acceptance_len: float
+    acceptance_len_rtol: float
+    expected_per_pos_acceptance_rates: tuple[float, ...]
+    per_pos_rtol: float
+    quantization: str | None = None
+    parallel_drafting: bool | None = None
 
 
-def compute_spec_decode_stats(
-    metrics,
-) -> dict:
+DFLASH_CONFIG = SpeculatorTestConfig(
+    model_path="nm-testing/dflash-qwen3-8b-speculators",
+    method="dflash",
+    display_name="DFlash",
+    expected_gsm8k_accuracy=0.885,
+    accuracy_rtol=0.03,
+    expected_acceptance_len=3.45,
+    acceptance_len_rtol=0.15,
+    expected_per_pos_acceptance_rates=(0.795, 0.611, 0.429, 0.282),
+    per_pos_rtol=0.15,
+    quantization="fp8",
+)
+
+PEAGLE_CONFIG = SpeculatorTestConfig(
+    model_path="nm-testing/qwen3-8b-peagle-speculators",
+    method="eagle3",
+    display_name="PEagle",
+    expected_gsm8k_accuracy=0.88,
+    accuracy_rtol=0.05,
+    expected_acceptance_len=2.27,
+    acceptance_len_rtol=0.20,
+    expected_per_pos_acceptance_rates=(0.66, 0.36, 0.18, 0.09),
+    per_pos_rtol=0.20,
+    parallel_drafting=True,
+)
+
+SPECULATOR_CONFIGS = [
+    pytest.param(DFLASH_CONFIG, id="dflash"),
+    pytest.param(PEAGLE_CONFIG, id="peagle"),
+]
+
+
+def compute_spec_decode_stats(metrics) -> dict:
     """Extract all spec-decode metrics and compute derived stats."""
     name2metric = {m.name: m for m in metrics}
 
@@ -67,25 +102,26 @@ def print_spec_decode_stats(stats: dict) -> None:
     print("===============================\n")
 
 
-def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
+@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
+def test_speculators_model(vllm_runner, example_prompts, monkeypatch, config):
     """
-    Test DFlash speculators model properly initializes speculative decoding.
+    Test speculators model properly initializes speculative decoding.
 
     Verifies:
     1. Speculative config is automatically initialized from speculators config
-    2. Method is detected as 'dflash'
-    3. The draft model path is correctly set
-    4. Speculative tokens count is valid (num_speculative_tokens=8)
-    5. Text generation works with speculative decoding enabled
+    2. Method is detected correctly
+    3. parallel_drafting is set correctly (if applicable)
+    4. The draft model path is correctly set
+    5. Speculative tokens count is valid
+    6. Text generation works with speculative decoding enabled
     """
     monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
 
-    with vllm_runner(
-        MODEL_PATH,
-        dtype=torch.bfloat16,
-        enforce_eager=True,
-        quantization="fp8",
-    ) as vllm_model:
+    runner_kwargs = dict(dtype=torch.bfloat16, enforce_eager=True)
+    if config.quantization:
+        runner_kwargs["quantization"] = config.quantization
+
+    with vllm_runner(config.model_path, **runner_kwargs) as vllm_model:
         vllm_config = vllm_model.llm.llm_engine.vllm_config
 
         assert isinstance(vllm_config.speculative_config, SpeculativeConfig), (
@@ -93,40 +129,43 @@ def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
         )
 
         spec_config = vllm_config.speculative_config
-        assert spec_config.method == "dflash", (
-            f"Expected method='dflash', got '{spec_config.method}'"
+        assert spec_config.method == config.method, (
+            f"Expected method='{config.method}', got '{spec_config.method}'"
         )
+        if config.parallel_drafting is not None:
+            assert spec_config.parallel_drafting is config.parallel_drafting, (
+                f"Expected parallel_drafting={config.parallel_drafting} "
+                f"for {config.display_name} model"
+            )
         assert spec_config.num_speculative_tokens > 0, (
             f"Expected positive speculative tokens, "
             f"got {spec_config.num_speculative_tokens}"
         )
-        assert spec_config.model == MODEL_PATH, (
-            f"Draft model should be {MODEL_PATH}, got {spec_config.model}"
+        assert spec_config.model == config.model_path, (
+            f"Draft model should be {config.model_path}, got {spec_config.model}"
         )
 
         vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20)
-        assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}"
+        assert vllm_outputs, (
+            f"No outputs generated for speculators model {config.model_path}"
+        )
 
 
 @pytest.mark.slow_test
 @large_gpu_mark(min_gb=40)
-def test_dflash_speculators_correctness(monkeypatch):
+@pytest.mark.parametrize("config", SPECULATOR_CONFIGS)
+def test_speculators_correctness(monkeypatch, config):
     """
-    E2E correctness test for DFlash via the speculators auto-detect path.
+    E2E correctness test via the speculators auto-detect path.
 
     Evaluates GSM8k accuracy to ensure the speculators-format model produces
     correct outputs, and checks that acceptance length does not collapse under
     batched inference (lm-eval style).
-
-    Observed per-position acceptance rates on GSM8K (1319 prompts):
-        pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282,
-        pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023
-    Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128)
     """
     monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
 
     spec_llm = LLM(
-        model=MODEL_PATH,
+        model=config.model_path,
         trust_remote_code=True,
         max_model_len=4096,
         max_num_seqs=128,
@@ -137,7 +176,7 @@ def test_dflash_speculators_correctness(monkeypatch):
 
     results = evaluate_gsm8k_offline(spec_llm)
     accuracy = results["accuracy"]
-    accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL)
+    accuracy_threshold = config.expected_gsm8k_accuracy * (1 - config.accuracy_rtol)
     assert accuracy >= accuracy_threshold, (
         f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
     )
@@ -147,19 +186,18 @@ def test_dflash_speculators_correctness(monkeypatch):
     print_spec_decode_stats(stats)
 
     acceptance_len = stats["acceptance_len"]
-    al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL)
+    al_threshold = config.expected_acceptance_len * (1 - config.acceptance_len_rtol)
     assert acceptance_len >= al_threshold, (
-        f"DFlash speculators acceptance length too low: "
+        f"{config.display_name} speculators acceptance length too low: "
         f"{acceptance_len:.2f} < {al_threshold:.2f}"
     )
 
-    # Check per-position acceptance rates for the first few positions.
     per_pos_rates = stats["per_pos_acceptance_rates"]
-    for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES):
+    for i, expected_rate in enumerate(config.expected_per_pos_acceptance_rates):
         assert i < len(per_pos_rates), (
             f"Missing per-position acceptance rate for position {i}"
         )
-        threshold = expected_rate * (1 - PER_POS_RTOL)
+        threshold = expected_rate * (1 - config.per_pos_rtol)
         assert per_pos_rates[i] >= threshold, (
             f"Per-position acceptance rate at pos {i} too low: "
             f"{per_pos_rates[i]:.4f} < {threshold:.4f} "

vllm/model_executor/models/registry.py

@@ -591,6 +591,8 @@ _SPECULATIVE_DECODING_MODELS = {
     "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),
     "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),
     "DFlashDraftModel": ("qwen3_dflash", "DFlashQwen3ForCausalLM"),
+    "PEagleDraftModel": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
+    "PeagleLlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "Eagle3MiniMaxM2ForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
     "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

vllm/transformers_utils/configs/speculators/algos.py

@@ -43,6 +43,37 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
         ]
 
 
+@register_speculator("peagle")
+def update_peagle(config_dict: dict, pre_trained_config: dict) -> None:
+    """
+    Apply PEagle (Parallel Eagle) specific configuration transformations to
+    the `dict` used to construct the Transformers PreTrainedConfig.
+
+    PEagle specific fields:
+    - draft_vocab_size: Size of the draft model's vocabulary
+    - target_hidden_size: Hidden size of the target model
+    - norm_before_residual: Whether to apply norm before residual connection
+    - norm_before_fc: Whether to apply RMSNorm before the fc projection
+    - mask_token_id (required): Token ID used for parallel drafting mask
+        placeholders, mapped to pard_token for the proposer
+    - eagle_aux_hidden_state_layer_ids: Layer indices from the target model
+        whose intermediate hidden states are used as auxiliary inputs
+    """
+    pre_trained_config["architectures"] = ["PeagleLlamaForCausalLM"]
+    pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
+    if config_dict.get("target_hidden_size") is not None:
+        pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
+    pre_trained_config["norm_before_residual"] = config_dict.get(
+        "norm_before_residual", False
+    )
+    pre_trained_config["norm_before_fc"] = config_dict.get("norm_before_fc", False)
+    pre_trained_config["pard_token"] = config_dict["mask_token_id"]
+    if config_dict.get("eagle_aux_hidden_state_layer_ids"):
+        pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
+            "eagle_aux_hidden_state_layer_ids"
+        ]
+
+
 @register_speculator("dflash")
 def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
     """

vllm/transformers_utils/configs/speculators/base.py

@@ -131,7 +131,10 @@ class SpeculatorsConfig(PretrainedConfig):
             )
 
         # Build base vLLM speculative configuration
-        return {
+        result = {
             "method": config_dict.get("speculators_model_type"),
             "num_speculative_tokens": num_speculative_tokens,
         }
+        if result["method"] == "peagle":
+            result.update({"method": "eagle3", "parallel_drafting": True})
+        return result