@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f1d53d150c

Fork 0

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, 05:30 UTC

[Multimodal][Speculative Decoding]Eagle Eagle3 mm support, enablement on qwen2.5vl (#22872)

Signed-off-by: Junhong <liujunhong11@huawei.com> Signed-off-by: Junhong Liu <98734602+LJH-LBJ@users.noreply.github.com> Co-authored-by: Junhong <liujunhong11@huawei.com> Co-authored-by: LJH-LBJ <98734602+LJH-LBJ@users.noreply.github.com>

WeiQing Chen committed Sep 27, 2025 at 11:35 UTC f1d53d150c5cd9c7d94db296793fc25f955ea8a9

8 files changed +210 -45

tests/models/registry.py

@@ -651,6 +651,9 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

651 "MiMoMTPModel": _HfExamplesInfo("XiaomiMiMo/MiMo-7B-RL",

652 trust_remote_code=True,

653 speculative_model="XiaomiMiMo/MiMo-7B-RL"),

654 + "Eagle3Qwen2_5vlForCausalLM": _HfExamplesInfo(

655 + "Qwen/Qwen2.5-VL-7B-Instruct",

656 + speculative_model="Rayzl/qwen2.5-vl-7b-eagle3-sgl"),

657 "Qwen3NextMTP": _HfExamplesInfo("Qwen/Qwen3-Next-80B-A3B-Instruct",

658 min_transformers_version="4.56.3"),

659 }

tests/v1/e2e/test_spec_decode.py

+7 -2

@@ -129,6 +129,11 @@ def test_ngram_correctness(

129 ["model_setup", "mm_enabled"],

130 [

131 (("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), False),

132 + pytest.param(("eagle3", "Qwen/Qwen2.5-VL-7B-Instruct",

133 + "Rayzl/qwen2.5-vl-7b-eagle3-sgl", 1),

134 + False,

135 + marks=pytest.mark.skip(reason="Skipping due to its " \

136 + "head_dim not being a a multiple of 32")),

137 (("eagle", "meta-llama/Llama-3.1-8B-Instruct",

138 "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", 1), False),

139 (("eagle3", "meta-llama/Llama-3.1-8B-Instruct",

@@ -145,8 +150,8 @@ def test_ngram_correctness(

150 "eagle618/eagle-deepseek-v3-random", 1), False),

151 ],

152 ids=[

148 - "qwen3_eagle3", "llama3_eagle", "llama3_eagle3", "llama4_eagle",

149 - "llama4_eagle_mm", "deepseek_eagle"

153 + "qwen3_eagle3", "qwen2_5_vl_eagle3", "llama3_eagle", "llama3_eagle3",

154 + "llama4_eagle", "llama4_eagle_mm", "deepseek_eagle"

155 ])

156 @pytest.mark.parametrize("attn_backend",

157 get_attn_backend_list_based_on_platform())

vllm/benchmarks/datasets.py

+80

@@ -1450,6 +1450,13 @@ def get_samples(args, tokenizer) -> list[SampleRequest]:

1450 ):

1451 dataset_class = MLPerfDataset

1452 args.hf_split = "train"

1453 + elif (

1454 + args.dataset_path in MMStarDataset.SUPPORTED_DATASET_PATHS

1455 + or args.hf_name in MMStarDataset.SUPPORTED_DATASET_PATHS

1456 + ):

1457 + dataset_class = MMStarDataset

1458 + args.hf_split = "val"

1459 + args.hf_subset = None

1460 else:

1461 supported_datasets = set([

1462 dataset_name for cls in HuggingFaceDataset.__subclasses__()

@@ -2721,3 +2728,76 @@ class PrefixRepetitionRandomDataset(BenchmarkDataset):

2728

2729 random.shuffle(requests)

2730 return requests

2731 +

2732 +

2733 +# -----------------------------------------------------------------------------

2734 +# MMStar Dataset Implementation

2735 +# -----------------------------------------------------------------------------

2736 +

2737 +

2738 +class MMStarDataset(HuggingFaceDataset):

2739 + """

2740 + Lin-Chen/MMStar: https://huggingface.co/datasets/Lin-Chen/MMStar

2741 + refer to: https://github.com/sgl-project/SpecForge/pull/106

2742 + """

2743 + DEFAULT_OUTPUT_LEN = 128

2744 + SUPPORTED_DATASET_PATHS = {"Lin-Chen/MMStar"}

2745 + IS_MULTIMODAL = True

2746 +

2747 + def sample(

2748 + self,

2749 + tokenizer: PreTrainedTokenizerBase,

2750 + num_requests: int,

2751 + output_len: Optional[int] = None,

2752 + enable_multimodal_chat: bool = False,

2753 + request_id_prefix: str = "",

2754 + no_oversample: bool = False,

2755 + **kwargs,

2756 + ) -> list[SampleRequest]:

2757 + # If --hf-output-len is not set, use the default output length.

2758 + output_len = (output_len

2759 + if output_len is not None else self.DEFAULT_OUTPUT_LEN)

2760 + sampled_requests: list[SampleRequest] = []

2761 +

2762 + for ind, item in enumerate(self.data):

2763 + if len(sampled_requests) >= num_requests:

2764 + break

2765 + # Split the question text from options

2766 + # (keep only the part before "Options:").

2767 + full_q: str = item.get("question", "")

2768 + question_text = full_q.split("Options:", 1)[0].strip()

2769 +

2770 + # Multimodal image content.

2771 + mm_content = process_image(item["image"])

2772 +

2773 + # Compute prompt token length (note: this is plain text length

2774 + # if enable_multimodal_chat is False).

2775 + prompt_len = len(tokenizer(question_text).input_ids)

2776 +

2777 + if enable_multimodal_chat:

2778 + # If multimodal content should be embedded in the chat message,

2779 + # convert to [{"role":"user","content":[...]}]

2780 + prompt = self.apply_multimodal_chat_transformation(

2781 + question_text, mm_content

2782 + )

2783 + mm_for_request = None # Already embedded in chat content.

2784 + else:

2785 + # Default: prompt is plain text,

2786 + # image is in mm_content for the bench to assemble.

2787 + prompt = question_text

2788 + mm_for_request = mm_content

2789 +

2790 + sampled_requests.append(

2791 + SampleRequest(

2792 + prompt=prompt,

2793 + prompt_len=prompt_len,

2794 + expected_output_len=output_len,

2795 + multi_modal_data=mm_for_request,

2796 + request_id=request_id_prefix + str(ind),

2797 + )

2798 + )

2799 +

2800 + self.maybe_oversample_requests(

2801 + sampled_requests, num_requests, request_id_prefix, no_oversample

2802 + )

2803 + return sampled_requests

vllm/model_executor/models/llama_eagle3.py

+19 -8

@@ -8,7 +8,6 @@ import torch

8 import torch.nn as nn

9 from transformers import LlamaConfig

11 -from vllm.compilation.decorators import support_torch_compile

11 from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.layernorm import RMSNorm

@@ -19,6 +18,7 @@ from vllm.model_executor.layers.quantization.base_config import (

18 from vllm.model_executor.layers.vocab_parallel_embedding import (

19 DEFAULT_VOCAB_PADDING_SIZE, ParallelLMHead, VocabParallelEmbedding)

20 from vllm.model_executor.model_loader.weight_utils import default_weight_loader

21 +from vllm.model_executor.models.interfaces import MultiModalEmbeddings

22 from vllm.model_executor.models.llama import (LlamaDecoderLayer,

23 LlamaForCausalLM)

@@ -102,7 +102,6 @@ class LlamaDecoderLayer(LlamaDecoderLayer):

102 return hidden_states, residual

103

104

105 -@support_torch_compile

105 class LlamaModel(nn.Module):

106

107 def __init__(

@@ -145,13 +144,21 @@ class LlamaModel(nn.Module):

144 eps=self.config.rms_norm_eps,

145 )

146

147 + def get_input_embeddings(

148 + self,

149 + input_ids: torch.Tensor,

150 + ) -> torch.Tensor:

151 + return self.embed_tokens(input_ids)

152 +

153 def forward(

154 self,

155 input_ids: torch.Tensor,

156 positions: torch.Tensor,

157 hidden_states: torch.Tensor,

158 + input_embeds: Optional[torch.Tensor] = None,

159 ) -> tuple[torch.Tensor, torch.Tensor]:

154 - input_embeds = self.embed_tokens(input_ids)

160 + if input_embeds is None:

161 + input_embeds = self.get_input_embeddings(input_ids)

162 assert hidden_states.shape[-1] == input_embeds.shape[-1]

163

164 residual = None

@@ -239,11 +246,7 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

246 hidden_states: torch.Tensor,

247 inputs_embeds: Optional[torch.Tensor] = None,

248 ) -> tuple[torch.Tensor, torch.Tensor]:

242 - if inputs_embeds is not None:

243 - raise NotImplementedError(

244 - f"{type(self).__name__} does not support multimodal inputs yet."

245 - )

246 - return self.model(input_ids, positions, hidden_states)

249 + return self.model(input_ids, positions, hidden_states, inputs_embeds)

250

251 def compute_logits(

252 self,

@@ -299,3 +302,11 @@ class Eagle3LlamaForCausalLM(LlamaForCausalLM):

302 skip_substrs=skip_substrs,

303 )

304 loader.load_weights(model_weights.items())

305 +

306 + def get_input_embeddings(

307 + self,

308 + input_ids: torch.Tensor,

309 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

310 + ) -> torch.Tensor:

311 + inputs_embeds = self.model.get_input_embeddings(input_ids)

312 + return inputs_embeds

vllm/model_executor/models/qwen2_5_vl.py

+9 -2

@@ -68,7 +68,7 @@ from vllm.transformers_utils.config import uses_mrope

68 from vllm.utils import is_pin_memory_available

69 from vllm.utils.tensor_schema import TensorSchema, TensorShape

71 -from .interfaces import (MultiModalEmbeddings, SupportsLoRA,

71 +from .interfaces import (MultiModalEmbeddings, SupportsEagle3, SupportsLoRA,

72 SupportsMultiModal, SupportsMultiModalPruning,

73 SupportsPP, SupportsQuant)

74 from .qwen2_vl import Qwen2VLDummyInputsBuilder as Qwen2_5_VLDummyInputsBuilder

@@ -965,7 +965,7 @@ class Qwen2_5_VLMultiModalProcessor(Qwen2VLMultiModalProcessor):

965 dummy_inputs=Qwen2_5_VLDummyInputsBuilder)

966 class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

967 SupportsLoRA, SupportsPP,

968 - SupportsQuant,

968 + SupportsQuant, SupportsEagle3,

969 SupportsMultiModalPruning):

970

971 packed_modules_mapping = {

@@ -1028,6 +1028,13 @@ class Qwen2_5_VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1028 self.make_empty_intermediate_tensors = (

1029 self.language_model.make_empty_intermediate_tensors)

1030

1031 + def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:

1032 + self.language_model.model.aux_hidden_state_layers = layers

1033 +

1034 + def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:

1035 + num_layers = len(self.language_model.model.layers)

1036 + return (2, num_layers // 2, num_layers - 3)

1037 +

1038 def _validate_and_reshape_mm_tensor(self, mm_input: object,

1039 name: str) -> torch.Tensor:

1040 if not isinstance(mm_input, (torch.Tensor, list)):

vllm/model_executor/models/registry.py

@@ -286,6 +286,7 @@ _SPECULATIVE_DECODING_MODELS = {

286 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),

287 "Eagle3LlamaForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

288 "LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

289 + "Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),

290 "EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),

291 "DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),

292 "ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),

vllm/v1/spec_decode/eagle.py

+79 -29

@@ -80,9 +80,17 @@ class EagleProposer:

80 self.input_ids = torch.zeros(self.max_num_tokens,

81 dtype=torch.int32,

82 device=device)

83 - self.positions = torch.zeros(self.max_num_tokens,

84 - dtype=torch.int64,

85 - device=device)

83 + self.uses_mrope = self.vllm_config.model_config.uses_mrope

84 + if self.uses_mrope:

85 + # M-RoPE need (3, max_num_tokens)

86 + self.mrope_positions = torch.zeros((3, self.max_num_tokens),

87 + dtype=torch.int64,

88 + device=device)

89 + else:

90 + # RoPE need (max_num_tokens,)

91 + self.positions = torch.zeros(self.max_num_tokens,

92 + dtype=torch.int64,

93 + device=device)

94 self.hidden_states = torch.zeros(

95 (self.max_num_tokens, self.hidden_size),

96 dtype=self.dtype,

@@ -143,11 +151,22 @@ class EagleProposer:

151 dtype=torch.int32,

152 ).repeat(max_batch_size, 1)

153

154 + def _get_positions(self, num_tokens: int):

155 + if self.uses_mrope:

156 + return self.mrope_positions[:, :num_tokens]

157 + return self.positions[:num_tokens]

158 +

159 + def _set_positions(self, num_tokens: int, positions: torch.Tensor):

160 + if self.uses_mrope:

161 + self.mrope_positions[:, :num_tokens] = positions

162 + else:

163 + self.positions[:num_tokens] = positions

164 +

165 def propose(

166 self,

167 # [num_tokens]

168 target_token_ids: torch.Tensor,

150 - # [num_tokens]

169 + # [num_tokens] or [3, num_tokens] when M-RoPE is enabled

170 target_positions: torch.Tensor,

171 # [num_tokens, hidden_size]

172 target_hidden_states: torch.Tensor,

@@ -198,7 +217,7 @@ class EagleProposer:

217 else:

218 num_input_tokens = num_tokens

219 # copy inputs to buffer for cudagraph

201 - self.positions[:num_tokens] = target_positions

220 + self._set_positions(num_tokens, target_positions)

221 self.hidden_states[:num_tokens] = target_hidden_states

222 if self.is_multimodal_model:

223 input_ids = self.input_ids[:num_tokens]

@@ -218,7 +237,7 @@ class EagleProposer:

237 num_tokens=num_input_tokens):

238 ret_hidden_states = self.model(

239 input_ids=input_ids,

221 - positions=self.positions[:num_input_tokens],

240 + positions=self._get_positions(num_input_tokens),

241 hidden_states=self.hidden_states[:num_input_tokens],

242 inputs_embeds=inputs_embeds,

243 )

@@ -235,7 +254,10 @@ class EagleProposer:

254 draft_token_ids = logits.argmax(dim=-1)

255 return draft_token_ids.view(-1, 1)

256

238 - positions = target_positions[last_token_indices]

257 + if self.uses_mrope:

258 + positions = target_positions[:, last_token_indices]

259 + else:

260 + positions = target_positions[last_token_indices]

261 if self.method in ("deepseek_mtp", "ernie_mtp", "longcat_flash_mtp"):

262 hidden_states = self.hidden_states[last_token_indices]

263 else:

@@ -282,25 +304,34 @@ class EagleProposer:

304 # cast to int32 is crucial when eagle model is compiled.

305 # tensor.argmax() returns int64 by default.

306 input_ids = draft_token_ids_list[-1].int()

285 - positions += 1

286 -

287 - # NOTE(woosuk): We should handle the case where the draft model

288 - # generates tokens beyond the max model length. Since it is complex

289 - # to remove such requests from the batch, we keep them in the batch

290 - # but adjust the position ids and slot mappings to avoid the

291 - # out-of-range access during the model execution. The draft tokens

292 - # generated with this adjustment should be ignored.

293 - exceeds_max_model_len = positions >= self.max_model_len

294 - # Mask out the position ids that exceed the max model length.

295 - # Otherwise, we may get out-of-range error in RoPE.

296 - clamped_positions = torch.where(exceeds_max_model_len, 0,

297 - positions)

307 + if self.uses_mrope:

308 + positions += 1

309 + # NOTE(woosuk): We should handle the case where the draft model

310 + # generates tokens beyond the max model length.

311 + # Since it is complex to remove such requests from the batch,

312 + # we keep them in the batch but adjust the position ids

313 + # and slot mappings to avoid the

314 + # out-of-range access during the model execution.

315 + # The draft tokens generated with this adjustment

316 + # should be ignored.

317 + exceeds_max_model_len = positions[0] >= self.max_model_len

318 + # Mask out the position ids that exceed the max model length.

319 + # Otherwise, we may get out-of-range error in RoPE.

320 + clamped_positions = torch.where\

321 + (exceeds_max_model_len.unsqueeze(0), \

322 + torch.zeros_like(positions), positions)

323 + else:

324 + positions += 1

325 + exceeds_max_model_len = positions >= self.max_model_len

326 + clamped_positions = torch.where(exceeds_max_model_len, 0,

327 + positions)

328

329 # Increment the sequence lengths.

330 common_attn_metadata.seq_lens += 1

331 common_attn_metadata.seq_lens_cpu += 1

332 # For the requests that exceed the max model length, we set the

333 # sequence length to 1 to minimize their overheads in attention.

334 +

335 common_attn_metadata.seq_lens.masked_fill_(exceeds_max_model_len,

336 1)

337

@@ -308,13 +339,22 @@ class EagleProposer:

339 common_attn_metadata.seq_lens_cpu - 1

340

341 # Compute the slot mapping.

311 - block_numbers = clamped_positions // self.block_size

342 + if self.uses_mrope:

343 + # all dimensions of positions are the same

344 + block_numbers = clamped_positions[0] // self.block_size

345 + else:

346 + block_numbers = clamped_positions // self.block_size

347 block_ids = common_attn_metadata.block_table_tensor.gather(

348 dim=1, index=block_numbers.view(-1, 1))

349 block_ids = block_ids.view(-1)

315 - common_attn_metadata.slot_mapping = (

316 - block_ids * self.block_size +

317 - clamped_positions % self.block_size)

350 + if self.uses_mrope:

351 + common_attn_metadata.slot_mapping = (

352 + block_ids * self.block_size +

353 + clamped_positions[0] % self.block_size)

354 + else:

355 + common_attn_metadata.slot_mapping = (

356 + block_ids * self.block_size +

357 + clamped_positions % self.block_size)

358 # Mask out the slot mappings that exceed the max model length.

359 # Otherwise, the KV cache will be inadvertently updated with the

360 # padding tokens.

@@ -330,7 +370,7 @@ class EagleProposer:

370

371 # copy inputs to buffer for cudagraph

372 self.input_ids[:batch_size] = input_ids

333 - self.positions[:batch_size] = clamped_positions

373 + self._set_positions(batch_size, clamped_positions)

374 self.hidden_states[:batch_size] = hidden_states

375 if self.is_multimodal_model:

376 inputs_embeds = self.model.get_input_embeddings(input_ids)

@@ -347,7 +387,7 @@ class EagleProposer:

387 num_tokens=input_batch_size):

388 ret_hidden_states = self.model(

389 input_ids=input_ids,

350 - positions=self.positions[:input_batch_size],

390 + positions=self._get_positions(input_batch_size),

391 hidden_states=self.hidden_states[:input_batch_size],

392 inputs_embeds=inputs_embeds,

393 )

@@ -787,6 +827,11 @@ class EagleProposer:

827

828 return spec_common_attn_metadata, token_indices

829

830 + def get_model_name(self, model: nn.Module) -> str:

831 + if hasattr(model, 'module'): # multi-GPU

832 + model = model.module

833 + return model.__class__.__name__

834 +

835 def load_model(self, target_model: nn.Module) -> None:

836 draft_model_config = \

837 self.vllm_config.speculative_config.draft_model_config

@@ -820,8 +865,13 @@ class EagleProposer:

865

866 if supports_multimodal(target_model):

867 # handle multimodality

823 - self.model.config.image_token_index = (

824 - target_model.config.image_token_index)

868 + if (self.get_model_name(target_model) ==

869 + "Qwen2_5_VLForConditionalGeneration"):

870 + self.model.config.image_token_index = (

871 + target_model.config.image_token_id)

872 + else:

873 + self.model.config.image_token_index = (

874 + target_model.config.image_token_index)

875 target_language_model = target_model.get_language_model()

876 else:

877 target_language_model = target_model

@@ -892,7 +942,7 @@ class EagleProposer:

942

943 self.model(

944 input_ids=input_ids,

895 - positions=self.positions[:num_tokens],

945 + positions=self._get_positions(num_tokens),

946 hidden_states=self.hidden_states[:num_tokens],

947 inputs_embeds=inputs_embeds,

948 )

vllm/v1/worker/gpu_model_runner.py

+12 -4

@@ -442,6 +442,16 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

442 device="cpu",

443 pin_memory=self.pin_memory)

444

445 + def _get_positions(self, num_tokens: Any):

446 + if isinstance(num_tokens, int):

447 + if self.uses_mrope:

448 + return self.mrope_positions.gpu[:, :num_tokens]

449 + return self.positions.gpu[:num_tokens]

450 + else:

451 + if self.uses_mrope:

452 + return self.mrope_positions.gpu[:, num_tokens]

453 + return self.positions.gpu[num_tokens]

454 +

455 def _make_buffer(self,

456 *size: Union[int, torch.SymInt],

457 dtype: torch.dtype,

@@ -2544,8 +2554,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2554 token_indices_to_sample = None

2555 # input_ids can be None for multimodal models.

2556 target_token_ids = self.input_ids.gpu[:num_scheduled_tokens]

2547 - # TODO(woosuk): Support M-RoPE.

2548 - target_positions = self.positions.gpu[:num_scheduled_tokens]

2557 + target_positions = self._get_positions(num_scheduled_tokens)

2558 if self.use_aux_hidden_state_outputs:

2559 assert aux_hidden_states is not None

2560 target_hidden_states = torch.cat(

@@ -2570,8 +2579,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

2579 valid_sampled_tokens_count)

2580

2581 target_token_ids = self.input_ids.gpu[token_indices]

2573 - # TODO(woosuk): Support M-RoPE.

2574 - target_positions = self.positions.gpu[token_indices]

2582 + target_positions = self._get_positions(token_indices)

2583 if self.use_aux_hidden_state_outputs:

2584 assert aux_hidden_states is not None

2585 target_hidden_states = torch.cat(