[BugFix] eagle draft max position embeddings (#49343)
Signed-off-by: JaredforReal <w13431838023@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Jared Wen committed
Jul 29, 2026 at 16:24 UTC
32e657e6894a8709d60f8edeffb5a2b1b1cb59a3
4 files changed
+288
-4
tests/config/test_speculative_draft_max_position_embeddings.py
new
+115
@@ -0,0 +1,115 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the EAGLE draft ``max_position_embeddings`` override (#48894).
4
+
5
+EAGLE drafts share the target's positional space, but some draft
6
+checkpoints (e.g. ``yuhuili/EAGLE3-LLaMA3.1-Instruct-8B``) ship a
7
+``max_position_embeddings`` (2048) far smaller than the target's context.
8
+That value sizes the draft's rotary ``cos_sin_cache`` while the proposer
9
+feeds positions up to the target's ``max_model_len``, so the cache gather
10
+goes out of bounds — a device-side assert under torch.compile and silent
11
+garbage reads in eager mode. ``SpeculativeConfig`` must raise the draft's
12
+value to the target's ``max_model_len``, with a log, for the eagle/eagle3
13
+methods only.
14
+"""
15
+
16
+import logging
17
+
18
+import pytest
19
+from transformers import PretrainedConfig
20
+
21
+from vllm.config.model import ModelConfig
22
+from vllm.config.parallel import ParallelConfig
23
+from vllm.config.speculative import SpeculativeConfig
24
+
25
+# All repos are public; only config/tokenizer-config files are fetched.
26
+EAGLE3_DRAFT = "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B" # max_position_embeddings=2048
27
+LLAMA3_TARGET = "unsloth/Meta-Llama-3.1-8B-Instruct" # max_position_embeddings=131072
28
+AR_MODEL = "JackFram/llama-68m" # max_position_embeddings=2048
29
+
30
+_LOGGER = "vllm.config.speculative"
31
+_OVERRIDE_MSG = "Overriding draft model max_position_embeddings"
32
+
33
+
34
+@pytest.fixture
35
+def vllm_caplog(caplog: pytest.LogCaptureFixture, monkeypatch: pytest.MonkeyPatch):
36
+ """Make caplog see vLLM logger records (vLLM sets propagate=False)."""
37
+ monkeypatch.setattr(logging.getLogger("vllm"), "propagate", True)
38
+ with caplog.at_level(logging.INFO, logger=_LOGGER):
39
+ yield caplog
40
+
41
+
42
+def _override_logged(caplog: pytest.LogCaptureFixture) -> bool:
43
+ return any(_OVERRIDE_MSG in record.getMessage() for record in caplog.records)
44
+
45
+
46
+@pytest.mark.cpu_test
47
+def test_override_raises_smaller_value(vllm_caplog: pytest.LogCaptureFixture):
48
+ hf_config = PretrainedConfig(max_position_embeddings=2048)
49
+ SpeculativeConfig._maybe_override_draft_max_position_embeddings(
50
+ hf_config, target_max_model_len=8192
51
+ )
52
+ assert hf_config.max_position_embeddings == 8192
53
+ assert _override_logged(vllm_caplog)
54
+
55
+
56
+@pytest.mark.cpu_test
57
+def test_override_keeps_sufficient_value(vllm_caplog: pytest.LogCaptureFixture):
58
+ hf_config = PretrainedConfig(max_position_embeddings=8192)
59
+ SpeculativeConfig._maybe_override_draft_max_position_embeddings(
60
+ hf_config, target_max_model_len=8192
61
+ )
62
+ assert hf_config.max_position_embeddings == 8192
63
+ assert not _override_logged(vllm_caplog)
64
+
65
+
66
+@pytest.mark.cpu_test
67
+def test_override_ignores_missing_attribute(vllm_caplog: pytest.LogCaptureFixture):
68
+ hf_config = PretrainedConfig()
69
+ hf_config.__dict__.pop("max_position_embeddings", None)
70
+ SpeculativeConfig._maybe_override_draft_max_position_embeddings(
71
+ hf_config, target_max_model_len=8192
72
+ )
73
+ assert not hasattr(hf_config, "max_position_embeddings")
74
+ assert not _override_logged(vllm_caplog)
75
+
76
+
77
+@pytest.mark.cpu_test
78
+@pytest.mark.parametrize("method", ["eagle", "eagle3"])
79
+def test_eagle_draft_inherits_target_max_model_len(
80
+ method: str, vllm_caplog: pytest.LogCaptureFixture
81
+):
82
+ target_model_config = ModelConfig(LLAMA3_TARGET)
83
+ assert target_model_config.max_model_len > 2048
84
+ speculative_config = SpeculativeConfig(
85
+ target_model_config=target_model_config,
86
+ target_parallel_config=ParallelConfig(),
87
+ model=EAGLE3_DRAFT,
88
+ method=method,
89
+ num_speculative_tokens=3,
90
+ )
91
+ draft_hf_config = speculative_config.draft_model_config.hf_config
92
+ assert draft_hf_config.max_position_embeddings == target_model_config.max_model_len
93
+ assert _override_logged(vllm_caplog)
94
+
95
+
96
+@pytest.mark.cpu_test
97
+def test_independent_draft_model_keeps_its_own_limit(
98
+ vllm_caplog: pytest.LogCaptureFixture,
99
+):
100
+ """An independent AR draft may genuinely have a smaller context than the
101
+ target; its max_position_embeddings must not be resized."""
102
+ target_model_config = ModelConfig(
103
+ AR_MODEL, hf_overrides={"max_position_embeddings": 8192}
104
+ )
105
+ assert target_model_config.max_model_len == 8192
106
+ speculative_config = SpeculativeConfig(
107
+ target_model_config=target_model_config,
108
+ target_parallel_config=ParallelConfig(),
109
+ model=AR_MODEL,
110
+ method="draft_model",
111
+ num_speculative_tokens=3,
112
+ )
113
+ draft_hf_config = speculative_config.draft_model_config.hf_config
114
+ assert draft_hf_config.max_position_embeddings == 2048
115
+ assert not _override_logged(vllm_caplog)
tests/v1/spec_decode/test_llm_base_proposer.py
new
+112
@@ -0,0 +1,112 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for SpecDecodeBaseProposer.initialize_attn_backend.
4
+
5
+Block tables are stored at kernel-block granularity, so the proposer's
6
+``block_size`` (used for slot-mapping math) must be the kernel block size,
7
+not the KV cache manager's block size — the two differ when manager blocks
8
+are split for the attention kernel. The value must also be deterministic:
9
+``_draft_attn_layer_names`` is a set, whose iteration order varies across
10
+processes, so anything derived from iteration order must not leak into
11
+``block_size``.
12
+"""
13
+
14
+from types import SimpleNamespace
15
+
16
+import pytest
17
+
18
+import vllm.v1.spec_decode.llm_base_proposer as llm_base_proposer
19
+from vllm.v1.spec_decode.eagle import EagleProposer
20
+
21
+SCHEDULER_BLOCK_SIZE = 256
22
+KERNEL_BLOCK_SIZE = 64
23
+
24
+
25
+class _FakeAttentionGroup:
26
+ def __init__(self, backend, layer_names, kv_cache_spec, kv_cache_group_id):
27
+ self.backend = backend
28
+ self.layer_names = list(layer_names)
29
+ self.kv_cache_spec = kv_cache_spec
30
+ self.kv_cache_group_id = kv_cache_group_id
31
+ self.kernel_block_size = None
32
+
33
+ def create_metadata_builders(self, vllm_config, device, kernel_block_size=None):
34
+ self.kernel_block_size = kernel_block_size
35
+
36
+ def get_metadata_builder(self):
37
+ return SimpleNamespace(kv_cache_spec=self.kv_cache_spec)
38
+
39
+
40
+def _make_proposer(
41
+ monkeypatch: pytest.MonkeyPatch, layer_names: set[str]
42
+) -> EagleProposer:
43
+ fake_layers = {}
44
+ for name in layer_names:
45
+ backend = SimpleNamespace(full_cls_name=lambda: "FakeBackend")
46
+ fake_layers[name] = SimpleNamespace(
47
+ get_attn_backend=lambda backend=backend: backend
48
+ )
49
+ monkeypatch.setattr(
50
+ llm_base_proposer, "get_layers_from_vllm_config", lambda *a, **k: fake_layers
51
+ )
52
+ monkeypatch.setattr(llm_base_proposer, "AttentionGroup", _FakeAttentionGroup)
53
+
54
+ proposer = EagleProposer.__new__(EagleProposer)
55
+ proposer.vllm_config = None
56
+ proposer.device = None
57
+ proposer._draft_attn_layer_names = set(layer_names)
58
+ proposer.kv_cache_gid = -1
59
+ proposer.draft_attn_groups = []
60
+ proposer.block_size = -1
61
+ return proposer
62
+
63
+
64
+def _make_kv_cache_config(layer_names: set[str]) -> SimpleNamespace:
65
+ spec = SimpleNamespace(block_size=SCHEDULER_BLOCK_SIZE)
66
+ group = SimpleNamespace(layer_names=list(layer_names), kv_cache_spec=spec)
67
+ return SimpleNamespace(kv_cache_groups=[group])
68
+
69
+
70
+def test_block_size_uses_kernel_block_size(monkeypatch: pytest.MonkeyPatch):
71
+ """The proposer's slot-mapping math runs against the kernel-granularity
72
+ block table, so block_size must come from kernel_block_sizes."""
73
+ layer_names = {"draft.0.self_attn.attn"}
74
+ proposer = _make_proposer(monkeypatch, layer_names)
75
+
76
+ proposer.initialize_attn_backend(
77
+ _make_kv_cache_config(layer_names),
78
+ kernel_block_sizes=[KERNEL_BLOCK_SIZE],
79
+ )
80
+
81
+ assert proposer.block_size == KERNEL_BLOCK_SIZE
82
+ assert proposer.block_size != SCHEDULER_BLOCK_SIZE
83
+ # The metadata builder keeps receiving the kernel block size as well.
84
+ assert proposer.draft_attn_groups[0].kernel_block_size == KERNEL_BLOCK_SIZE
85
+
86
+
87
+def test_block_size_falls_back_to_kv_cache_spec(monkeypatch: pytest.MonkeyPatch):
88
+ layer_names = {"draft.0.self_attn.attn"}
89
+ proposer = _make_proposer(monkeypatch, layer_names)
90
+
91
+ proposer.initialize_attn_backend(
92
+ _make_kv_cache_config(layer_names), kernel_block_sizes=None
93
+ )
94
+
95
+ assert proposer.block_size == SCHEDULER_BLOCK_SIZE
96
+
97
+
98
+def test_draft_layer_iteration_is_deterministic(monkeypatch: pytest.MonkeyPatch):
99
+ """_draft_attn_layer_names is a set; the attention groups built from it
100
+ must not depend on its (process-random) iteration order."""
101
+ layer_names = {"draft.c.attn", "draft.a.attn", "draft.b.attn"}
102
+ expected_order = sorted(layer_names)
103
+
104
+ for insertion_order in (expected_order, expected_order[::-1]):
105
+ proposer = _make_proposer(monkeypatch, set(insertion_order))
106
+ proposer.initialize_attn_backend(
107
+ _make_kv_cache_config(set(insertion_order)),
108
+ kernel_block_sizes=[KERNEL_BLOCK_SIZE],
109
+ )
110
+ assert len(proposer.draft_attn_groups) == 1
111
+ assert proposer.draft_attn_groups[0].layer_names == expected_order
112
+ assert proposer.block_size == KERNEL_BLOCK_SIZE
vllm/config/speculative.py
+43
@@ -910,6 +910,15 @@ class SpeculativeConfig:
910
f"Unsupported speculative method: '{self.method}'"
911
)
912
913
+ if self.method in ("eagle", "eagle3"):
914
+ # EAGLE drafts share the target's positional space; a
915
+ # draft checkpoint with a smaller max_position_embeddings
916
+ # than the target under-sizes its rotary cache (#48894).
917
+ SpeculativeConfig._maybe_override_draft_max_position_embeddings(
918
+ self.draft_model_config.hf_config,
919
+ self.target_model_config.max_model_len,
920
+ )
921
+
922
# Replace hf_config for EAGLE draft_model
923
if self.method in ("eagle", "eagle3", "dflash"):
924
from vllm.transformers_utils.configs.eagle import EAGLEConfig
@@ -1130,6 +1139,40 @@ class SpeculativeConfig:
1139
)
1140
return result
1141
1142
+ @staticmethod
1143
+ def _maybe_override_draft_max_position_embeddings(
1144
+ draft_hf_config: PretrainedConfig,
1145
+ target_max_model_len: int,
1146
+ ) -> None:
1147
+ """Raise an EAGLE draft's max_position_embeddings up to the target's.
1148
+
1149
+ The proposer feeds the draft positions up to the target's
1150
+ max_model_len, while max_position_embeddings sizes the draft's
1151
+ rotary cos_sin_cache. A smaller checkpoint value (e.g. 2048 for
1152
+ yuhuili/EAGLE3-LLaMA3.1-Instruct-8B) makes that cache gather go
1153
+ out of bounds (#48894).
1154
+
1155
+ Args:
1156
+ draft_hf_config: The draft model's HF config, mutated in place.
1157
+ target_max_model_len: The target model's max_model_len.
1158
+ """
1159
+ draft_max_position_embeddings = getattr(
1160
+ draft_hf_config, "max_position_embeddings", None
1161
+ )
1162
+ if (
1163
+ draft_max_position_embeddings is None
1164
+ or draft_max_position_embeddings >= target_max_model_len
1165
+ ):
1166
+ return
1167
+ logger.info(
1168
+ "Overriding draft model max_position_embeddings from %d to the "
1169
+ "target model's max_model_len (%d); EAGLE drafts share the "
1170
+ "target's positional space.",
1171
+ draft_max_position_embeddings,
1172
+ target_max_model_len,
1173
+ )
1174
+ draft_hf_config.max_position_embeddings = target_max_model_len
1175
+
1176
@staticmethod
1177
def _verify_and_get_draft_tp(
1178
target_parallel_config: ParallelConfig,
vllm/v1/spec_decode/llm_base_proposer.py
+18
-4
@@ -1740,7 +1740,10 @@ class SpecDecodeBaseProposer:
1740
1741
attention_groups: dict[tuple[str, str], AttentionGroup] = {}
1742
if kv_cache_spec is not None:
1743
- for layer_name in self._draft_attn_layer_names:
1743
+ # _draft_attn_layer_names is a set; iterate in sorted order so
1744
+ # that attention_groups (and anything derived from its first
1745
+ # element) is deterministic across processes.
1746
+ for layer_name in sorted(self._draft_attn_layer_names):
1747
attn_backend = all_attn_layers[layer_name].get_attn_backend()
1748
backend_key = attn_backend.full_cls_name()
1749
if backend_key not in attention_groups:
@@ -1772,9 +1775,20 @@ class SpecDecodeBaseProposer:
1775
attention_groups[backend_key].layer_names.append(layer_name)
1776
1777
self.draft_attn_groups = list(attention_groups.values())
1775
- self.block_size = (
1776
- self.draft_attn_groups[0].get_metadata_builder().kv_cache_spec.block_size
1777
- )
1778
+ if kernel_block_sizes is not None and 0 <= self.kv_cache_gid < len(
1779
+ kernel_block_sizes
1780
+ ):
1781
+ # Slot mappings are computed against the block table, which is
1782
+ # stored at kernel-block granularity. Use the kernel block size
1783
+ # rather than the KV cache manager's block size; the two differ
1784
+ # when manager blocks are split for the attention kernel.
1785
+ self.block_size = kernel_block_sizes[self.kv_cache_gid]
1786
+ else:
1787
+ self.block_size = (
1788
+ self.draft_attn_groups[0]
1789
+ .get_metadata_builder()
1790
+ .kv_cache_spec.block_size
1791
+ )
1792
logger.debug("Using block size %d for drafting layers", self.block_size)
1793
1794
def _determine_batch_execution_and_padding(