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 14, 22:35 UTC
[Meta] Official Eagle mm support, first enablement on llama4 (#20788)
Signed-off-by: morgendave <morgendave@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.me>
zhiweiz committed
Jul 31, 2025 at 10:35 UTC
9e0726e5bfd201fa2c9209e3997d24c72ecc3b13
8 files changed
+206
-37
examples/offline_inference/spec_decode.py
+59
-7
@@ -13,6 +13,38 @@ except ImportError:
13
from argparse import ArgumentParser as FlexibleArgumentParser
14
15
16
+QUESTION = "What is the content of each image?"
17
+IMAGE_URLS = [
18
+ "https://upload.wikimedia.org/wikipedia/commons/d/da/2015_Kaczka_krzy%C5%BCowka_w_wodzie_%28samiec%29.jpg",
19
+ "https://upload.wikimedia.org/wikipedia/commons/7/77/002_The_lion_king_Snyggve_in_the_Serengeti_National_Park_Photo_by_Giles_Laurent.jpg",
20
+ "https://upload.wikimedia.org/wikipedia/commons/2/26/Ultramarine_Flycatcher_%28Ficedula_superciliaris%29_Naggar%2C_Himachal_Pradesh%2C_2013_%28cropped%29.JPG",
21
+ "https://upload.wikimedia.org/wikipedia/commons/thumb/e/e5/Anim1754_-_Flickr_-_NOAA_Photo_Library_%281%29.jpg/2560px-Anim1754_-_Flickr_-_NOAA_Photo_Library_%281%29.jpg",
22
+ "https://upload.wikimedia.org/wikipedia/commons/d/d4/Starfish%2C_Caswell_Bay_-_geograph.org.uk_-_409413.jpg",
23
+ "https://upload.wikimedia.org/wikipedia/commons/6/69/Grapevinesnail_01.jpg",
24
+ "https://upload.wikimedia.org/wikipedia/commons/thumb/0/0b/Texas_invasive_Musk_Thistle_1.jpg/1920px-Texas_invasive_Musk_Thistle_1.jpg",
25
+ "https://upload.wikimedia.org/wikipedia/commons/thumb/7/7a/Huskiesatrest.jpg/2880px-Huskiesatrest.jpg",
26
+ "https://upload.wikimedia.org/wikipedia/commons/thumb/6/68/Orange_tabby_cat_sitting_on_fallen_leaves-Hisashi-01A.jpg/1920px-Orange_tabby_cat_sitting_on_fallen_leaves-Hisashi-01A.jpg",
27
+ "https://upload.wikimedia.org/wikipedia/commons/3/30/George_the_amazing_guinea_pig.jpg",
28
+ "https://upload.wikimedia.org/wikipedia/commons/thumb/1/1f/Oryctolagus_cuniculus_Rcdo.jpg/1920px-Oryctolagus_cuniculus_Rcdo.jpg",
29
+ "https://upload.wikimedia.org/wikipedia/commons/9/98/Horse-and-pony.jpg",
30
+]
31
+
32
+
33
+def get_custom_mm_prompts(num_prompts):
34
+ prompts = []
35
+ for url in IMAGE_URLS:
36
+ prompts.append(
37
+ [
38
+ {"type": "image_url", "image_url": {"url": url}},
39
+ {"type": "text", "text": QUESTION},
40
+ ]
41
+ )
42
+ if num_prompts > len(IMAGE_URLS):
43
+ prompts = prompts * (num_prompts // len(IMAGE_URLS) + 1)
44
+
45
+ return [[{"role": "user", "content": prompt}] for prompt in prompts[:num_prompts]]
46
+
47
+
48
def parse_args():
49
parser = FlexibleArgumentParser()
50
add_dataset_parser(parser)
@@ -35,6 +67,7 @@ def parse_args():
67
parser.add_argument("--output-len", type=int, default=256)
68
parser.add_argument("--model-dir", type=str, default=None)
69
parser.add_argument("--eagle-dir", type=str, default=None)
70
+ parser.add_argument("--custom-mm-prompts", action="store_true")
71
return parser.parse_args()
72
73
@@ -44,14 +77,26 @@ def main():
77
78
model_dir = args.model_dir
79
if args.model_dir is None:
80
+ if args.custom_mm_prompts:
81
+ raise ValueError(
82
+ "custom_mm_prompts requires mm based models"
83
+ "default llama3.1-8b-instruct is not mm based"
84
+ "please specify model_dir to give a mm based model"
85
+ )
86
model_dir = "meta-llama/Llama-3.1-8B-Instruct"
87
tokenizer = AutoTokenizer.from_pretrained(model_dir)
49
-
50
- prompts = get_samples(args, tokenizer)
51
- # add_special_tokens is False to avoid adding bos twice when using chat templates
52
- prompt_ids = [
53
- tokenizer.encode(prompt.prompt, add_special_tokens=False) for prompt in prompts
54
- ]
88
+ args.custom_skip_chat_template = True
89
+
90
+ if not args.custom_mm_prompts:
91
+ prompts = get_samples(args, tokenizer)
92
+ # add_special_tokens is False to avoid adding bos twice
93
+ # when using chat templates
94
+ prompt_ids = [
95
+ tokenizer.encode(prompt.prompt, add_special_tokens=False)
96
+ for prompt in prompts
97
+ ]
98
+ else:
99
+ prompts = get_custom_mm_prompts(args.num_prompts)
100
101
if args.method == "eagle" or args.method == "eagle3":
102
eagle_dir = args.eagle_dir
@@ -85,10 +130,17 @@ def main():
130
speculative_config=speculative_config,
131
disable_log_stats=False,
132
max_model_len=16384,
133
+ limit_mm_per_prompt={"image": 5},
134
+ disable_chunked_mm_input=True,
135
)
136
137
sampling_params = SamplingParams(temperature=args.temp, max_tokens=args.output_len)
91
- outputs = llm.generate(prompt_token_ids=prompt_ids, sampling_params=sampling_params)
138
+ if not args.custom_mm_prompts:
139
+ outputs = llm.generate(
140
+ prompt_token_ids=prompt_ids, sampling_params=sampling_params
141
+ )
142
+ else:
143
+ outputs = llm.chat(prompts, sampling_params=sampling_params)
144
145
# print the generated text
146
if args.print_output:
tests/v1/e2e/test_spec_decode.py
+45
-16
@@ -3,29 +3,34 @@
3
from __future__ import annotations
4
5
import random
6
-from typing import Any
6
+from typing import Any, Union
7
8
import pytest
9
import torch
10
11
from vllm import LLM, SamplingParams
12
+from vllm.assets.base import VLLM_S3_BUCKET_URL
13
+from vllm.assets.image import VLM_IMAGES_DIR
14
from vllm.distributed import cleanup_dist_env_and_memory
15
16
15
-@pytest.fixture
16
-def test_prompts():
17
+def get_test_prompts(mm_enabled: bool):
18
prompt_types = ["repeat", "sentence"]
19
+ if mm_enabled:
20
+ prompt_types.append("mm")
21
num_prompts = 100
22
prompts = []
23
24
random.seed(0)
25
random_prompt_type_choices = random.choices(prompt_types, k=num_prompts)
26
+ print(f"Prompt types: {random_prompt_type_choices}")
27
28
# Generate a mixed batch of prompts, some of which can be easily
29
# predicted by n-gram matching and some which likely cannot.
30
for kind in random_prompt_type_choices:
31
word_choices = ["test", "temp", "hello", "where"]
32
word = random.choice(word_choices)
33
+ prompt: Union[str, list[dict[str, Any]]] = ""
34
if kind == "repeat":
35
prompt = f"""
36
please repeat the word '{word}' 10 times.
@@ -38,6 +43,21 @@ def test_prompts():
43
uses the word {word} at least once.
44
give no other output than that simple sentence without quotes.
45
"""
46
+ elif kind == "mm":
47
+ placeholders = [{
48
+ "type": "image_url",
49
+ "image_url": {
50
+ "url":
51
+ f"{VLLM_S3_BUCKET_URL}/{VLM_IMAGES_DIR}/stop_sign.jpg"
52
+ },
53
+ }]
54
+ prompt = [
55
+ *placeholders,
56
+ {
57
+ "type": "text",
58
+ "text": "The meaning of the image is"
59
+ },
60
+ ]
61
else:
62
raise ValueError(f"Unknown prompt type: {kind}")
63
prompts.append([{"role": "user", "content": prompt}])
@@ -57,7 +77,6 @@ def model_name():
77
78
def test_ngram_correctness(
79
monkeypatch: pytest.MonkeyPatch,
60
- test_prompts: list[list[dict[str, Any]]],
80
sampling_config: SamplingParams,
81
model_name: str,
82
):
@@ -67,6 +86,7 @@ def test_ngram_correctness(
86
'''
87
with monkeypatch.context() as m:
88
m.setenv("VLLM_USE_V1", "1")
89
+ test_prompts = get_test_prompts(mm_enabled=False)
90
91
ref_llm = LLM(model=model_name, max_model_len=1024)
92
ref_outputs = ref_llm.chat(test_prompts, sampling_config)
@@ -103,23 +123,32 @@ def test_ngram_correctness(
123
cleanup_dist_env_and_memory()
124
125
106
-@pytest.mark.parametrize("model_setup", [
107
- ("eagle", "meta-llama/Llama-3.1-8B-Instruct",
108
- "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", 1),
109
- ("eagle3", "meta-llama/Llama-3.1-8B-Instruct",
110
- "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", 1),
111
- pytest.param(
112
- ("eagle", "meta-llama/Llama-4-Scout-17B-16E-Instruct",
113
- "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", 4),
114
- marks=pytest.mark.skip(reason="Skipping due to CI OOM issues")),
115
-],
116
- ids=["llama3_eagle", "llama3_eagle3", "llama4_eagle"])
126
+@pytest.mark.parametrize(
127
+ ["model_setup", "mm_enabled"], [
128
+ (("eagle", "meta-llama/Llama-3.1-8B-Instruct",
129
+ "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", 1), False),
130
+ (("eagle3", "meta-llama/Llama-3.1-8B-Instruct",
131
+ "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", 1), False),
132
+ pytest.param(
133
+ ("eagle", "meta-llama/Llama-4-Scout-17B-16E-Instruct",
134
+ "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", 4),
135
+ False,
136
+ marks=pytest.mark.skip(reason="Skipping due to CI OOM issues")),
137
+ pytest.param(
138
+ ("eagle", "meta-llama/Llama-4-Scout-17B-16E-Instruct",
139
+ "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", 4),
140
+ True,
141
+ marks=pytest.mark.skip(reason="Skipping due to CI OOM issues")),
142
+ ],
143
+ ids=["llama3_eagle", "llama3_eagle3", "llama4_eagle", "llama4_eagle_mm"])
144
def test_eagle_correctness(
145
monkeypatch: pytest.MonkeyPatch,
119
- test_prompts: list[list[dict[str, Any]]],
146
sampling_config: SamplingParams,
147
model_setup: tuple[str, str, str, int],
148
+ mm_enabled: bool,
149
):
150
+ # Generate test prompts inside the function instead of using fixture
151
+ test_prompts = get_test_prompts(mm_enabled)
152
'''
153
Compare the outputs of a original LLM and a speculative LLM
154
should be the same when using eagle speculative decoding.
vllm/model_executor/models/llama4.py
+1
@@ -256,6 +256,7 @@ class Llama4DecoderLayer(nn.Module):
256
super().__init__()
257
258
self.layer_idx = extract_layer_index(prefix)
259
+ self.global_layer = config.no_rope_layers[self.layer_idx] == 0
260
self.hidden_size = config.hidden_size
261
rope_theta = config.rope_theta
262
rope_scaling = config.rope_scaling
vllm/model_executor/models/llama4_eagle.py
+31
-4
@@ -37,8 +37,9 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
37
from vllm.model_executor.models.llama4 import (Llama4DecoderLayer,
38
Llama4ForCausalLM)
39
from vllm.model_executor.models.utils import extract_layer_index
40
+from vllm.multimodal.inputs import NestedTensors
41
41
-from .utils import AutoWeightsLoader, maybe_prefix
42
+from .utils import AutoWeightsLoader, maybe_prefix, merge_multimodal_embeddings
43
44
logger = init_logger(__name__)
45
@@ -78,15 +79,23 @@ class LlamaModel(nn.Module):
79
self.norm = RMSNorm(self.config.hidden_size,
80
eps=self.config.rms_norm_eps)
81
82
+ def get_input_embeddings(
83
+ self,
84
+ input_ids: torch.Tensor,
85
+ ) -> torch.Tensor:
86
+ return self.embed_tokens(input_ids)
87
+
88
def forward(
89
self,
90
input_ids: Optional[torch.Tensor],
91
positions: torch.Tensor,
92
hidden_states: torch.Tensor,
93
+ inputs_embeds: Optional[torch.Tensor] = None,
94
) -> tuple[torch.Tensor, torch.Tensor]:
87
- input_embeds = self.embed_tokens(input_ids)
95
+ if inputs_embeds is None:
96
+ inputs_embeds = self.get_input_embeddings(input_ids)
97
hidden_states = self.fc(
89
- torch.cat((input_embeds, hidden_states), dim=-1))
98
+ torch.cat((inputs_embeds, hidden_states), dim=-1))
99
residual = None
100
for layer in self.layers:
101
hidden_states, residual = layer(
@@ -190,8 +199,9 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
199
input_ids: torch.Tensor,
200
positions: torch.Tensor,
201
hidden_states: torch.Tensor,
202
+ inputs_embeds: Optional[torch.Tensor] = None,
203
) -> tuple[torch.Tensor, torch.Tensor]:
194
- return self.model(input_ids, positions, hidden_states)
204
+ return self.model(input_ids, positions, hidden_states, inputs_embeds)
205
206
def load_weights(self, weights: Iterable[tuple[str,
207
torch.Tensor]]) -> None:
@@ -212,3 +222,20 @@ class EagleLlama4ForCausalLM(Llama4ForCausalLM):
222
model_weights[name] = loaded_weight
223
224
loader.load_weights(model_weights.items())
225
+
226
+ def get_input_embeddings(
227
+ self,
228
+ input_ids: torch.Tensor,
229
+ multimodal_embeddings: Optional[NestedTensors] = None,
230
+ ) -> torch.Tensor:
231
+ inputs_embeds = self.model.get_input_embeddings(input_ids)
232
+
233
+ if multimodal_embeddings is not None:
234
+ inputs_embeds = merge_multimodal_embeddings(
235
+ input_ids,
236
+ inputs_embeds,
237
+ multimodal_embeddings,
238
+ self.config.image_token_index,
239
+ )
240
+
241
+ return inputs_embeds
vllm/model_executor/models/llama_eagle.py
+6
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from collections.abc import Iterable
5
+from typing import Optional
6
7
import torch
8
import torch.nn as nn
@@ -148,7 +149,12 @@ class EagleLlamaForCausalLM(LlamaForCausalLM):
149
input_ids: torch.Tensor,
150
positions: torch.Tensor,
151
hidden_states: torch.Tensor,
152
+ inputs_embeds: Optional[torch.Tensor] = None,
153
) -> tuple[torch.Tensor, torch.Tensor]:
154
+ if inputs_embeds is not None:
155
+ raise NotImplementedError(
156
+ f"{type(self).__name__} does not support multimodal inputs yet."
157
+ )
158
return self.model(input_ids, positions, hidden_states)
159
160
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
vllm/model_executor/models/llama_eagle3.py
+5
@@ -202,7 +202,12 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):
202
input_ids: torch.Tensor,
203
positions: torch.Tensor,
204
hidden_states: torch.Tensor,
205
+ inputs_embeds: Optional[torch.Tensor] = None,
206
) -> tuple[torch.Tensor, torch.Tensor]:
207
+ if inputs_embeds is not None:
208
+ raise NotImplementedError(
209
+ f"{type(self).__name__} does not support multimodal inputs yet."
210
+ )
211
return self.model(input_ids, positions, hidden_states)
212
213
def compute_logits(
vllm/v1/spec_decode/eagle.py
+50
-9
@@ -1,5 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+from typing import Optional
4
+
5
import numpy as np
6
import torch
7
import torch.nn as nn
@@ -51,6 +53,9 @@ class EagleProposer:
53
# hidden size (e.g., Llama 3.3 70B).
54
self.hidden_size = self.draft_model_config.get_hidden_size()
55
56
+ self.is_multimodal_model = vllm_config.model_config \
57
+ .is_multimodal_model
58
+
59
self.use_cuda_graph = (self.vllm_config.compilation_config.level
60
== CompilationLevel.PIECEWISE and
61
not self.vllm_config.model_config.enforce_eager)
@@ -76,6 +81,11 @@ class EagleProposer:
81
device=device,
82
dtype=torch.int32)
83
84
+ self.inputs_embeds = torch.zeros(
85
+ (self.max_num_tokens, self.hidden_size),
86
+ dtype=self.dtype,
87
+ device=device)
88
+
89
def propose(
90
self,
91
# [num_tokens]
@@ -88,6 +98,7 @@ class EagleProposer:
98
next_token_ids: torch.Tensor,
99
common_attn_metadata: CommonAttentionMetadata,
100
sampling_metadata: SamplingMetadata,
101
+ mm_embeds: Optional[list[torch.Tensor]] = None,
102
) -> torch.Tensor:
103
num_tokens = target_token_ids.shape[0]
104
batch_size = next_token_ids.shape[0]
@@ -128,14 +139,27 @@ class EagleProposer:
139
# copy inputs to buffer for cudagraph
140
self.positions[:num_tokens] = target_positions
141
self.hidden_states[:num_tokens] = target_hidden_states
142
+ if self.is_multimodal_model:
143
+ input_ids = self.input_ids[:num_tokens]
144
+ inputs_embeds = self.model.get_input_embeddings(
145
+ input_ids,
146
+ multimodal_embeddings=mm_embeds or None,
147
+ )
148
+ self.inputs_embeds[:num_tokens] = inputs_embeds
149
+ inputs_embeds = self.inputs_embeds[:num_input_tokens]
150
+ input_ids = None
151
+ else:
152
+ inputs_embeds = None
153
+ input_ids = self.input_ids[:num_input_tokens]
154
155
with set_forward_context(per_layer_attn_metadata,
156
self.vllm_config,
157
num_tokens=num_input_tokens):
158
ret_hidden_states = self.model(
136
- self.input_ids[:num_input_tokens],
137
- self.positions[:num_input_tokens],
138
- self.hidden_states[:num_input_tokens],
159
+ input_ids=input_ids,
160
+ positions=self.positions[:num_input_tokens],
161
+ hidden_states=self.hidden_states[:num_input_tokens],
162
+ inputs_embeds=inputs_embeds,
163
)
164
if self.method == "deepseek_mtp":
165
last_hidden_states = ret_hidden_states
@@ -218,15 +242,24 @@ class EagleProposer:
242
self.input_ids[:batch_size] = input_ids
243
self.positions[:batch_size] = clamped_positions
244
self.hidden_states[:batch_size] = hidden_states
245
+ if self.is_multimodal_model:
246
+ inputs_embeds = self.model.get_input_embeddings(input_ids)
247
+ self.inputs_embeds[:batch_size] = inputs_embeds
248
+ inputs_embeds = self.inputs_embeds[:input_batch_size]
249
+ input_ids = None
250
+ else:
251
+ inputs_embeds = None
252
+ input_ids = self.input_ids[:input_batch_size]
253
254
# Run the model.
255
with set_forward_context(per_layer_attn_metadata,
256
self.vllm_config,
257
num_tokens=input_batch_size):
258
last_hidden_states, hidden_states = self.model(
227
- self.input_ids[:input_batch_size],
228
- self.positions[:input_batch_size],
229
- self.hidden_states[:input_batch_size],
259
+ input_ids=input_ids,
260
+ positions=self.positions[:input_batch_size],
261
+ hidden_states=self.hidden_states[:input_batch_size],
262
+ inputs_embeds=inputs_embeds,
263
)
264
hidden_states = hidden_states[:batch_size]
265
logits = self.model.compute_logits(last_hidden_states[:batch_size],
@@ -391,10 +424,18 @@ class EagleProposer:
424
) -> None:
425
with set_forward_context(None, self.vllm_config,
426
num_tokens=num_tokens):
427
+ if self.is_multimodal_model:
428
+ input_ids = None
429
+ inputs_embeds = self.inputs_embeds[:num_tokens]
430
+ else:
431
+ input_ids = self.input_ids[:num_tokens]
432
+ inputs_embeds = None
433
+
434
self.model(
395
- self.input_ids[:num_tokens],
396
- self.positions[:num_tokens],
397
- self.hidden_states[:num_tokens],
435
+ input_ids=input_ids,
436
+ positions=self.positions[:num_tokens],
437
+ hidden_states=self.hidden_states[:num_tokens],
438
+ inputs_embeds=inputs_embeds,
439
)
440
441
def validate_same_kv_cache_group(self,
vllm/v1/worker/gpu_model_runner.py
+9
-1
@@ -1205,13 +1205,15 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1205
def _gather_mm_embeddings(
1206
self,
1207
scheduler_output: "SchedulerOutput",
1208
+ shift_computed_tokens: int = 0,
1209
) -> list[torch.Tensor]:
1210
mm_embeds: list[torch.Tensor] = []
1211
for req_id in self.input_batch.req_ids:
1212
num_scheduled_tokens = scheduler_output.num_scheduled_tokens[
1213
req_id]
1214
req_state = self.requests[req_id]
1214
- num_computed_tokens = req_state.num_computed_tokens
1215
+ num_computed_tokens = \
1216
+ req_state.num_computed_tokens + shift_computed_tokens
1217
mm_positions = req_state.mm_positions
1218
for i, pos_info in enumerate(mm_positions):
1219
start_pos = pos_info.offset
@@ -1858,6 +1860,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1860
[h[token_indices] for h in aux_hidden_states], dim=-1)
1861
else:
1862
target_hidden_states = hidden_states[token_indices]
1863
+ mm_embeds = None
1864
+ if self.is_multimodal_model:
1865
+ mm_embeds = self._gather_mm_embeddings(scheduler_output,
1866
+ shift_computed_tokens=1)
1867
+
1868
draft_token_ids = self.drafter.propose(
1869
target_token_ids=target_token_ids,
1870
target_positions=target_positions,
@@ -1865,6 +1872,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1872
next_token_ids=next_token_ids,
1873
sampling_metadata=sampling_metadata,
1874
common_attn_metadata=common_attn_metadata,
1875
+ mm_embeds=mm_embeds,
1876
)
1877
spec_token_ids = draft_token_ids.tolist()
1878
return spec_token_ids