[Bugfix][Spec Decode] Fix EAGLE and dense draft startup with EP (#56930)
Signed-off-by: Kyle Liang <kyleliang-nv@users.noreply.github.com> Co-authored-by: Kyle Liang <kyleliang-nv@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>
kyleliang-nv committed
Sep 16, 2026 at 17:03 UTC
c58ff535f4a2f30b61a86c846ac4adefe62abb8d
2 files changed
+122
-3
tests/test_config.py
+109
@@ -1520,6 +1520,115 @@ def test_draft_model_enables_async_scheduling_by_default():
1520
assert cfg.scheduler_config.async_scheduling is True
1521
1522
1523
+@pytest.mark.skip_global_cleanup
1524
+@pytest.mark.parametrize("tp_size", [1, 2])
1525
+@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"])
1526
+@pytest.mark.parametrize(
1527
+ ("method", "draft_is_moe"),
1528
+ [
1529
+ pytest.param("draft_model", False, id="dense-draft"),
1530
+ pytest.param("eagle", False, id="eagle"),
1531
+ pytest.param("eagle3", False, id="eagle3"),
1532
+ pytest.param("draft_model", True, id="moe-draft"),
1533
+ pytest.param("mtp", True, id="mtp"),
1534
+ pytest.param("dspark", True, id="dspark"),
1535
+ ],
1536
+)
1537
+def test_draft_inherits_ep_only_for_moe(
1538
+ tmp_path: Path,
1539
+ tp_size: int,
1540
+ target_ep: bool,
1541
+ method: str,
1542
+ draft_is_moe: bool,
1543
+):
1544
+ """Validate final draft configs without loading weights or mocking validation."""
1545
+ from transformers import LlamaConfig, MixtralConfig
1546
+
1547
+ from vllm.transformers_utils.configs.deepseek_v4 import DeepseekV4Config
1548
+
1549
+ common = dict(
1550
+ hidden_size=128,
1551
+ intermediate_size=256,
1552
+ num_hidden_layers=2,
1553
+ num_attention_heads=4,
1554
+ num_key_value_heads=2,
1555
+ vocab_size=128,
1556
+ max_position_embeddings=2048,
1557
+ )
1558
+ if method in ("mtp", "dspark"):
1559
+ target_hf_config = DeepseekV4Config(
1560
+ architectures=["DeepseekV4ForCausalLM"],
1561
+ n_routed_experts=4,
1562
+ num_experts_per_tok=2,
1563
+ num_nextn_predict_layers=1,
1564
+ compress_ratios=[1, 1],
1565
+ head_dim=32,
1566
+ **common,
1567
+ )
1568
+ else:
1569
+ target_hf_config = MixtralConfig(
1570
+ architectures=["MixtralForCausalLM"], num_local_experts=4, **common
1571
+ )
1572
+ target_path = tmp_path / "target"
1573
+ draft_path = tmp_path / "draft"
1574
+ _write_json(target_path / "config.json", target_hf_config.to_dict())
1575
+ _write_json(
1576
+ draft_path / "config.json",
1577
+ LlamaConfig(architectures=["LlamaForCausalLM"], **common).to_dict(),
1578
+ )
1579
+ target_model_config = ModelConfig(
1580
+ model=str(target_path), tokenizer_mode="skip", max_model_len=2048
1581
+ )
1582
+ target_parallel_config = ParallelConfig(
1583
+ tensor_parallel_size=tp_size,
1584
+ enable_expert_parallel=target_ep,
1585
+ distributed_executor_backend="mp",
1586
+ )
1587
+ target_model_config.verify_with_parallel_config(target_parallel_config)
1588
+ speculative_config = SpeculativeConfig(
1589
+ method=method,
1590
+ model=str(target_path if draft_is_moe else draft_path),
1591
+ num_speculative_tokens=1,
1592
+ target_model_config=target_model_config,
1593
+ target_parallel_config=target_parallel_config,
1594
+ )
1595
+
1596
+ assert speculative_config.method == method
1597
+ assert speculative_config.draft_model_config.is_moe is draft_is_moe
1598
+ assert speculative_config.draft_parallel_config.enable_expert_parallel is (
1599
+ target_ep and draft_is_moe
1600
+ )
1601
+ assert speculative_config.draft_parallel_config.tensor_parallel_size == tp_size
1602
+ assert target_parallel_config.enable_expert_parallel is target_ep
1603
+ assert target_parallel_config.tensor_parallel_size == tp_size
1604
+
1605
+
1606
+@pytest.mark.skip_global_cleanup
1607
+@pytest.mark.parametrize("target_ep", [False, True], ids=["ep-off", "ep-on"])
1608
+@pytest.mark.parametrize("pass_none", [False, True], ids=["omitted", "explicit-none"])
1609
+def test_draft_parallel_config_preserves_ep_without_model(
1610
+ target_ep: bool, pass_none: bool
1611
+):
1612
+ """Legacy callers without draft model information keep EP inheritance."""
1613
+ target_parallel_config = ParallelConfig(
1614
+ tensor_parallel_size=2,
1615
+ enable_expert_parallel=target_ep,
1616
+ distributed_executor_backend="mp",
1617
+ )
1618
+ if pass_none:
1619
+ draft_parallel_config = SpeculativeConfig.create_draft_parallel_config(
1620
+ target_parallel_config, 2, draft_model_config=None
1621
+ )
1622
+ else:
1623
+ draft_parallel_config = SpeculativeConfig.create_draft_parallel_config(
1624
+ target_parallel_config, 2
1625
+ )
1626
+
1627
+ assert draft_parallel_config.enable_expert_parallel is target_ep
1628
+ assert draft_parallel_config.tensor_parallel_size == 2
1629
+ assert target_parallel_config.enable_expert_parallel is target_ep
1630
+
1631
+
1632
@pytest.mark.parametrize(
1633
("method", "parallel_drafting", "expected_slots"),
1634
[
vllm/config/speculative.py
+13
-3
@@ -1545,9 +1545,13 @@ class SpeculativeConfig:
1545
)
1546
)
1547
1548
+ # Use the final draft config, after any architecture overrides,
1549
+ # so EP agrees with ModelConfig's expert-count validation.
1550
self.draft_parallel_config = (
1551
SpeculativeConfig.create_draft_parallel_config(
1550
- self.target_parallel_config, self.draft_tensor_parallel_size
1552
+ self.target_parallel_config,
1553
+ self.draft_tensor_parallel_size,
1554
+ draft_model_config=self.draft_model_config,
1555
)
1556
)
1557
@@ -1743,15 +1747,21 @@ class SpeculativeConfig:
1747
def create_draft_parallel_config(
1748
target_parallel_config: ParallelConfig,
1749
speculative_draft_tensor_parallel_size: int,
1750
+ draft_model_config: ModelConfig | None = None,
1751
) -> ParallelConfig:
1752
"""Create a parallel config for use by the draft worker.
1753
1749
- This is mostly a copy of the target parallel config, except the tp_size.
1754
+ Use the draft TP size and disable inherited EP for known dense drafts.
1755
+ Without a draft model config, preserve the previous EP inheritance.
1756
"""
1757
+ enable_ep = target_parallel_config.enable_expert_parallel
1758
+ if draft_model_config is not None:
1759
+ enable_ep = enable_ep and draft_model_config.is_moe
1760
+
1761
draft_parallel_config = ParallelConfig(
1762
pipeline_parallel_size=1,
1763
tensor_parallel_size=speculative_draft_tensor_parallel_size,
1754
- enable_expert_parallel=target_parallel_config.enable_expert_parallel,
1764
+ enable_expert_parallel=enable_ep,
1765
distributed_executor_backend=target_parallel_config.distributed_executor_backend,
1766
max_parallel_loading_workers=target_parallel_config.max_parallel_loading_workers,
1767
disable_custom_all_reduce=target_parallel_config.disable_custom_all_reduce,