[Bugfix][Spec Decode] Fix EAGLE drafter multimodal encoder cache misses (#46315)
Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Roger Wang <hey@rogerw.io>
Nick Hill committed
Jun 22, 2026 at 11:14 UTC
3e6529cc0e0f039bcc8c82d1403efa8355fd053d
7 files changed
+331
-18
tests/v1/core/test_scheduler.py
+32
@@ -4708,6 +4708,38 @@ def test_free_encoder_inputs_respects_unconfirmed_placeholders():
4708
assert manager.get_cached_input_ids(request) == set()
4709
4710
4711
+def test_free_encoder_inputs_defers_for_eagle_lookahead():
4712
+ """With EAGLE speculative decoding, the encoder input is retained one extra
4713
+ position so the drafter's +1 look-ahead mm-embedding gather (which reads one
4714
+ position past the target's computed range) still finds it cached. This is
4715
+ the primary mechanism that prevents the drafter "Encoder cache miss"; the
4716
+ worker-side token-embedding fallback is only a backstop."""
4717
+ scheduler = create_scheduler(model="llava-hf/llava-1.5-7b-hf")
4718
+ # create_scheduler only builds ngram spec configs; force the eagle path that
4719
+ # _free_encoder_inputs keys off (self.use_eagle).
4720
+ scheduler.use_eagle = True
4721
+ mm_positions = [[PlaceholderRange(offset=50, length=100)]]
4722
+ request = create_requests(
4723
+ num_requests=1,
4724
+ num_tokens=250,
4725
+ mm_positions=mm_positions,
4726
+ )[0]
4727
+ manager = scheduler.encoder_cache_manager
4728
+ manager.allocate(request, 0)
4729
+ mm_end = 150 # offset + length
4730
+
4731
+ # Confirmed progress reaches the range end: without spec decode this frees
4732
+ # (see test below), but the drafter's +1 look-ahead still needs it.
4733
+ request.num_computed_tokens = mm_end
4734
+ scheduler._free_encoder_inputs(request)
4735
+ assert manager.get_cached_input_ids(request) == {0}
4736
+
4737
+ # One position past the range end: the +1 look-ahead has now passed it.
4738
+ request.num_computed_tokens = mm_end + 1
4739
+ scheduler._free_encoder_inputs(request)
4740
+ assert manager.get_cached_input_ids(request) == set()
4741
+
4742
+
4743
def test_free_encoder_inputs_unchanged_without_spec_decode():
4744
"""Without speculative decoding, encoder inputs are freed as soon as
4745
num_computed_tokens passes the placeholder range, as before."""
tests/v1/worker/test_encoder_runner.py
new
+157
@@ -0,0 +1,157 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for EncoderRunner.gather_mm_embeddings (model runner V2).
4
+
5
+Covers the speculative-drafter encoder-cache handling: the drafter reads one
6
+position ahead of the target model (``draft_lookahead``). The +1 look-ahead
7
+feature past the processed boundary is used when its encoder output is present
8
+and tolerated (token-embedding fallback) when it is not, while a miss within
9
+the processed range still fails loudly.
10
+"""
11
+
12
+import numpy as np
13
+import pytest
14
+import torch
15
+
16
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
17
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
18
+from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
19
+
20
+pytestmark = pytest.mark.cpu_test
21
+
22
+HIDDEN = 4
23
+
24
+
25
+def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
26
+ return MultiModalFeatureSpec(
27
+ data=None,
28
+ modality="image",
29
+ identifier=identifier,
30
+ mm_position=PlaceholderRange(offset=offset, length=length),
31
+ )
32
+
33
+
34
+def _make_runner(
35
+ features: list[MultiModalFeatureSpec],
36
+ cached: list[MultiModalFeatureSpec],
37
+) -> EncoderRunner:
38
+ cache = EncoderCache()
39
+ cache.mm_features["req0"] = features
40
+ for f in cached:
41
+ length = f.mm_position.length
42
+ cache.encoder_outputs[f.identifier] = torch.arange(
43
+ length * HIDDEN, dtype=torch.float32
44
+ ).reshape(length, HIDDEN)
45
+ return EncoderRunner(
46
+ model=None, # unused by gather_mm_embeddings
47
+ max_num_tokens=64,
48
+ hidden_size=HIDDEN,
49
+ encoder_cache=cache,
50
+ dtype=torch.float32,
51
+ device=torch.device("cpu"),
52
+ )
53
+
54
+
55
+def _gather(runner: EncoderRunner, *, num_scheduled: int, draft_lookahead: int):
56
+ # Single prefilling request, computed_prefill=0, prefill_len large.
57
+ return runner.gather_mm_embeddings(
58
+ req_ids=["req0"],
59
+ total_num_scheduled_tokens=num_scheduled,
60
+ num_scheduled_tokens=np.array([num_scheduled]),
61
+ query_start_loc=np.array([0]),
62
+ prefill_lens=np.array([1000]),
63
+ computed_prefill_lens=np.array([0]),
64
+ draft_lookahead=draft_lookahead,
65
+ )
66
+
67
+
68
+def test_draft_lookahead_uses_boundary_feature_when_cached():
69
+ """The drafter's +1 look-ahead can reach the feature at offset ==
70
+ processed_end (the next chunk). When its encoder output is already cached
71
+ (the scheduler encoded it ahead), it is used for the look-ahead position
72
+ rather than ignored."""
73
+ f0 = _feature("h0", offset=0, length=8)
74
+ f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end
75
+ runner = _make_runner([f0, f1], cached=[f0, f1])
76
+
77
+ mm_embeds, is_mm_embed = _gather(runner, num_scheduled=8, draft_lookahead=1)
78
+
79
+ # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7.
80
+ assert len(mm_embeds) == 2
81
+ assert bool(is_mm_embed[7])
82
+ assert int(is_mm_embed.sum()) == 8
83
+
84
+
85
+def test_draft_lookahead_tolerates_missing_boundary_feature():
86
+ """When the +1 look-ahead feature past the processed boundary is not yet
87
+ encoded, fall back to the token embedding (the draft token is verified by
88
+ the target) instead of raising."""
89
+ f0 = _feature("h0", offset=0, length=8)
90
+ f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached
91
+ runner = _make_runner([f0, f1], cached=[f0])
92
+
93
+ mm_embeds, is_mm_embed = _gather(runner, num_scheduled=8, draft_lookahead=1)
94
+
95
+ # Only f0 is gathered; f1's boundary position falls back silently.
96
+ assert len(mm_embeds) == 1
97
+ assert not bool(is_mm_embed[7])
98
+ assert int(is_mm_embed.sum()) == 7
99
+
100
+
101
+def test_draft_lookahead_raises_on_interior_miss():
102
+ """A miss for a feature within the processed range (not the look-ahead
103
+ boundary) is a real invariant violation and must fail loudly, even on the
104
+ drafter path."""
105
+ f0 = _feature("h0", offset=0, length=8) # interior, within processed range
106
+ runner = _make_runner([f0], cached=[])
107
+
108
+ with pytest.raises(RuntimeError, match="Encoder cache miss"):
109
+ _gather(runner, num_scheduled=8, draft_lookahead=1)
110
+
111
+
112
+def test_target_path_raises_on_encoder_cache_miss():
113
+ """On the target path (no look-ahead) a miss is a real invariant
114
+ violation and must fail loudly."""
115
+ f0 = _feature("h0", offset=0, length=8)
116
+ runner = _make_runner([f0], cached=[])
117
+
118
+ with pytest.raises(RuntimeError, match="Encoder cache miss"):
119
+ _gather(runner, num_scheduled=8, draft_lookahead=0)
120
+
121
+
122
+@pytest.mark.parametrize("draft_lookahead", [0, 1])
123
+def test_multi_request_batch_gathers_per_request(draft_lookahead):
124
+ """Two prefilling requests in one batch: per-request query bounds must be
125
+ indexed by request, not applied as whole arrays."""
126
+ a0 = _feature("a0", offset=0, length=8)
127
+ b0 = _feature("b0", offset=0, length=8)
128
+ cache = EncoderCache()
129
+ cache.mm_features["req0"] = [a0]
130
+ cache.mm_features["req1"] = [b0]
131
+ for f in (a0, b0):
132
+ cache.encoder_outputs[f.identifier] = torch.arange(
133
+ f.mm_position.length * HIDDEN, dtype=torch.float32
134
+ ).reshape(f.mm_position.length, HIDDEN)
135
+ runner = EncoderRunner(
136
+ model=None,
137
+ max_num_tokens=64,
138
+ hidden_size=HIDDEN,
139
+ encoder_cache=cache,
140
+ dtype=torch.float32,
141
+ device=torch.device("cpu"),
142
+ )
143
+
144
+ mm_embeds, is_mm_embed = runner.gather_mm_embeddings(
145
+ req_ids=["req0", "req1"],
146
+ total_num_scheduled_tokens=16,
147
+ num_scheduled_tokens=np.array([8, 8]),
148
+ query_start_loc=np.array([0, 8]),
149
+ prefill_lens=np.array([1000, 1000]),
150
+ computed_prefill_lens=np.array([0, 0]),
151
+ draft_lookahead=draft_lookahead,
152
+ )
153
+
154
+ # Both requests contribute a feature; with the +1 skew each marks 7 of its
155
+ # 8 positions (the skew drops one), otherwise all 8.
156
+ assert len(mm_embeds) == 2
157
+ assert int(is_mm_embed.sum()) == (14 if draft_lookahead else 16)
tests/v1/worker/test_gpu_model_runner_mm_gather.py
new
+99
@@ -0,0 +1,99 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for GPUModelRunner._gather_mm_embeddings (model runner V1).
4
+
5
+Mirrors tests/v1/worker/test_encoder_runner.py (the V2 runner): the EAGLE/MTP
6
+drafter reads one position ahead of the target (shift_computed_tokens=1). The
7
++1 look-ahead feature past the processed boundary is used when its encoder
8
+output is present and tolerated (token-embedding fallback) when it is not,
9
+while a miss within the processed range still fails loudly.
10
+
11
+`_gather_mm_embeddings` only uses CPU-side state, so it is exercised against a
12
+lightweight stub for `self` instead of a full (CUDA-only) runner.
13
+"""
14
+
15
+from types import SimpleNamespace
16
+
17
+import pytest
18
+import torch
19
+
20
+from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
21
+from vllm.v1.worker.gpu_model_runner import GPUModelRunner
22
+
23
+pytestmark = pytest.mark.cpu_test
24
+
25
+HIDDEN = 4
26
+
27
+
28
+def _feature(identifier: str, offset: int, length: int) -> MultiModalFeatureSpec:
29
+ return MultiModalFeatureSpec(
30
+ data=None,
31
+ modality="image",
32
+ identifier=identifier,
33
+ mm_position=PlaceholderRange(offset=offset, length=length),
34
+ )
35
+
36
+
37
+def _gather(features, cached, *, num_scheduled, shift, num_computed=0):
38
+ encoder_cache = {
39
+ f.identifier: torch.arange(
40
+ f.mm_position.length * HIDDEN, dtype=torch.float32
41
+ ).reshape(f.mm_position.length, HIDDEN)
42
+ for f in cached
43
+ }
44
+ req_state = SimpleNamespace(num_computed_tokens=num_computed, mm_features=features)
45
+ runner = SimpleNamespace(
46
+ input_batch=SimpleNamespace(req_ids=["req0"]),
47
+ requests={"req0": req_state},
48
+ encoder_cache=encoder_cache,
49
+ is_multimodal_pruning_enabled=False,
50
+ uses_mrope=False,
51
+ )
52
+ scheduler_output = SimpleNamespace(
53
+ total_num_scheduled_tokens=num_scheduled,
54
+ num_scheduled_tokens={"req0": num_scheduled},
55
+ )
56
+ return GPUModelRunner._gather_mm_embeddings(
57
+ runner, scheduler_output, shift_computed_tokens=shift
58
+ )
59
+
60
+
61
+def test_draft_shift_uses_boundary_feature_when_cached():
62
+ """The drafter's +1 look-ahead reaches the feature at offset ==
63
+ processed_end; when it is already cached it is used for the look-ahead
64
+ position rather than ignored."""
65
+ f0 = _feature("h0", offset=0, length=8)
66
+ f1 = _feature("h1", offset=8, length=8) # starts exactly at processed_end
67
+ mm_embeds, is_mm_embed = _gather([f0, f1], [f0, f1], num_scheduled=8, shift=1)
68
+
69
+ # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7.
70
+ assert len(mm_embeds) == 2
71
+ assert bool(is_mm_embed[7])
72
+ assert int(is_mm_embed.sum()) == 8
73
+
74
+
75
+def test_draft_shift_tolerates_missing_boundary_feature():
76
+ """When the +1 look-ahead feature past the processed boundary is not yet
77
+ encoded, fall back to the token embedding instead of raising."""
78
+ f0 = _feature("h0", offset=0, length=8)
79
+ f1 = _feature("h1", offset=8, length=8) # boundary feature, not cached
80
+ mm_embeds, is_mm_embed = _gather([f0, f1], [f0], num_scheduled=8, shift=1)
81
+
82
+ assert len(mm_embeds) == 1 # only f0; f1's boundary position falls back
83
+ assert not bool(is_mm_embed[7])
84
+ assert int(is_mm_embed.sum()) == 7
85
+
86
+
87
+def test_draft_shift_raises_on_interior_miss():
88
+ """A miss for a feature within the processed range (not the look-ahead
89
+ boundary) is a real invariant violation, even on the drafter path."""
90
+ f0 = _feature("h0", offset=0, length=8) # interior, within processed range
91
+ with pytest.raises(RuntimeError, match="Encoder cache miss"):
92
+ _gather([f0], [], num_scheduled=8, shift=1)
93
+
94
+
95
+def test_target_path_raises_on_encoder_cache_miss():
96
+ """On the target path (no shift) a miss is a real invariant violation."""
97
+ f0 = _feature("h0", offset=0, length=8)
98
+ with pytest.raises(RuntimeError, match="Encoder cache miss"):
99
+ _gather([f0], [], num_scheduled=8, shift=0)
vllm/v1/core/sched/scheduler.py
+9
-5
@@ -1872,6 +1872,11 @@ class Scheduler(SchedulerInterface):
1872
if not cached_encoder_input_ids:
1873
return
1874
1875
+ # Defer the free by the drafter's look-ahead so an entry stays
1876
+ # referenced until the drafter's +1 read has also passed it, mirroring
1877
+ # the shift the encoder scheduling path applies.
1878
+ spec_lookahead = 1 if self.use_eagle else 0
1879
+
1880
# Here, we use list(set) to avoid modifying the set while iterating
1881
# over it.
1882
for input_id in list(cached_encoder_input_ids):
@@ -1884,13 +1889,12 @@ class Scheduler(SchedulerInterface):
1889
# KVs have been calculated and cached already.
1890
self.encoder_cache_manager.free_encoder_input(request, input_id)
1891
elif (
1887
- start_pos + num_tokens
1892
+ start_pos + num_tokens + spec_lookahead
1893
<= request.num_computed_tokens - request.num_output_placeholders
1894
):
1890
- # The encoder output is already processed and stored in the
1891
- # decoder's KV cache, and progress is far enough past the
1892
- # placeholder range that no pending draft-token rejection can
1893
- # roll num_computed_tokens back into it.
1895
+ # Processed, stored in the decoder KV cache, and far enough past
1896
+ # the placeholder range (plus the drafter's look-ahead) that no
1897
+ # rejection or drafter gather can reference it.
1898
self.encoder_cache_manager.free_encoder_input(request, input_id)
1899
1900
def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
vllm/v1/worker/gpu/mm/encoder_runner.py
+21
-10
@@ -68,15 +68,19 @@ class EncoderRunner:
68
query_start_loc: np.ndarray,
69
prefill_lens: np.ndarray,
70
computed_prefill_lens: np.ndarray,
71
+ draft_lookahead: int = 0,
72
) -> tuple[list[torch.Tensor], torch.Tensor]:
72
- is_prefilling = (computed_prefill_lens < prefill_lens).tolist()
73
- all_decode = not any(is_prefilling)
74
- if all_decode:
73
+ if draft_lookahead:
74
+ computed_prefill_lens = computed_prefill_lens + draft_lookahead
75
+
76
+ is_prefilling_np = computed_prefill_lens < prefill_lens
77
+ if not is_prefilling_np.any():
78
# All decode requests, so no need to gather any embeddings.
79
return [], torch.zeros(
80
total_num_scheduled_tokens, dtype=torch.bool, device=self.device
81
)
82
83
+ is_prefilling = is_prefilling_np.tolist()
84
query_start = computed_prefill_lens.tolist()
85
query_end = (computed_prefill_lens + num_scheduled_tokens).tolist()
86
@@ -89,11 +93,12 @@ class EncoderRunner:
93
# OPTIMIZATION: Skip decode requests.
94
continue
95
96
+ cur_query_start = query_start[i]
97
+ cur_query_end = query_end[i]
98
+
99
mm_features = self.encoder_cache.mm_features[req_id]
100
lo, hi = get_mm_features_in_window(
94
- mm_features,
95
- start=query_start[i],
96
- end=query_end[i],
101
+ mm_features, start=cur_query_start, end=cur_query_end
102
)
103
for idx in range(lo, hi):
104
mm_feature = mm_features[idx]
@@ -101,8 +106,8 @@ class EncoderRunner:
106
start_pos = pos_info.offset
107
num_encoder_tokens = pos_info.length
108
104
- start_idx = max(query_start[i] - start_pos, 0)
105
- end_idx = min(query_end[i] - start_pos, num_encoder_tokens)
109
+ start_idx = max(cur_query_start - start_pos, 0)
110
+ end_idx = min(cur_query_end - start_pos, num_encoder_tokens)
111
assert start_idx < end_idx
112
curr_embeds_start, curr_embeds_end = (
113
pos_info.get_embeds_indices_in_range(start_idx, end_idx)
@@ -114,7 +119,13 @@ class EncoderRunner:
119
120
mm_hash = mm_feature.identifier
121
encoder_output = self.encoder_cache.encoder_outputs.get(mm_hash, None)
117
- assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
122
+ if encoder_output is None:
123
+ # A feature starting at/after the processed boundary is only
124
+ # reached via the drafter's +1 look-ahead and might not be
125
+ # encoded yet; fall back to the token embedding for drafting.
126
+ if start_pos + draft_lookahead >= cur_query_end:
127
+ continue
128
+ raise RuntimeError(f"Encoder cache miss for {mm_hash}.")
129
130
if (is_embed := pos_info.is_embed) is not None:
131
is_embed = is_embed[start_idx:end_idx]
@@ -122,7 +133,7 @@ class EncoderRunner:
133
else:
134
mm_embeds_item = encoder_output[start_idx:end_idx]
135
125
- req_start_pos = query_start_loc[i] + start_pos - query_start[i]
136
+ req_start_pos = query_start_loc[i] + start_pos - cur_query_start
137
is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= (
138
True if is_embed is None else is_embed
139
)
vllm/v1/worker/gpu/model_runner.py
+3
-2
@@ -1409,8 +1409,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1409
input_batch.num_scheduled_tokens,
1410
input_batch.query_start_loc_np,
1411
input_batch.prefill_len_np,
1412
- # +1 to consider the skew in eagle
1413
- input_batch.num_computed_prefill_tokens_np + 1,
1412
+ input_batch.num_computed_prefill_tokens_np,
1413
+ # The EAGLE/MTP drafter reads one position ahead of the target.
1414
+ draft_lookahead=1,
1415
)
1416
1417
# Postprocess results and update request states.
vllm/v1/worker/gpu_model_runner.py
+10
-1
@@ -3152,7 +3152,16 @@ class GPUModelRunner(
3152
3153
mm_hash = mm_feature.identifier
3154
encoder_output = self.encoder_cache.get(mm_hash, None)
3155
- assert encoder_output is not None, f"Encoder cache miss for {mm_hash}."
3155
+ if encoder_output is None:
3156
+ # A feature starting at/after the processed boundary is only
3157
+ # reached via the drafter's +1 look-ahead and might not be
3158
+ # encoded yet; fall back to the token embedding for drafting.
3159
+ if (
3160
+ start_pos
3161
+ >= req_state.num_computed_tokens + num_scheduled_tokens
3162
+ ):
3163
+ continue
3164
+ raise RuntimeError(f"Encoder cache miss for {mm_hash}.")
3165
3166
if (is_embed := pos_info.is_embed) is not None:
3167
is_embed = is_embed[start_idx:end_idx]