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 4, 13:45 UTC
[FEATURE] Add EagleMistralForCausalLM (#41024)
Signed-off-by: juliendenize <julien.denize@mistral.ai>
Julien Denize committed
Apr 28, 2026 at 21:22 UTC
e9f8f31e9a4c31d6842ca1adffe2619ed204fafb
3 files changed
+172
tests/models/registry.py
+5
@@ -1464,6 +1464,11 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1464
speculative_model="yuhuili/EAGLE3-LLaMA3.1-Instruct-8B",
1465
tokenizer="MiniMaxAI/MiniMax-M2",
1466
),
1467
+ "EagleMistralForCausalLM": _HfExamplesInfo(
1468
+ "mistralai/Mistral-Medium-3.5-128B",
1469
+ speculative_model="mistralai/Mistral-Medium-3.5-128B-EAGLE",
1470
+ is_available_online=False,
1471
+ ),
1472
"EagleMistralLarge3ForCausalLM": _HfExamplesInfo(
1473
"mistralai/Mistral-Large-3-675B-Instruct-2512",
1474
speculative_model="mistralai/Mistral-Large-3-675B-Instruct-2512-Eagle",
vllm/model_executor/models/mistral_eagle.py
new
+166
@@ -0,0 +1,166 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from collections.abc import Iterable
5
+
6
+import torch
7
+import torch.nn as nn
8
+
9
+from vllm.compilation.decorators import support_torch_compile
10
+from vllm.config import VllmConfig
11
+from vllm.logger import init_logger
12
+from vllm.model_executor.layers.layernorm import RMSNorm
13
+from vllm.model_executor.layers.linear import RowParallelLinear
14
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
15
+from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
16
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
17
+from vllm.model_executor.models.interfaces import MultiModalEmbeddings
18
+from vllm.model_executor.models.llama import LlamaConfig
19
+from vllm.model_executor.models.mistral import (
20
+ MistralDecoderLayer,
21
+ MistralForCausalLM,
22
+ MistralModel,
23
+)
24
+from vllm.model_executor.models.utils import (
25
+ _merge_multimodal_embeddings,
26
+ get_draft_quant_config,
27
+ maybe_prefix,
28
+)
29
+
30
+logger = init_logger(__name__)
31
+
32
+
33
+class EagleMistralDecoderLayer(MistralDecoderLayer):
34
+ def __init__(
35
+ self,
36
+ vllm_config: VllmConfig,
37
+ prefix: str = "",
38
+ config: LlamaConfig | None = None,
39
+ ) -> None:
40
+ super().__init__(vllm_config, prefix=prefix, config=config)
41
+
42
+ def get_quant_config(self, vllm_config: VllmConfig) -> QuantizationConfig | None:
43
+ return get_draft_quant_config(vllm_config)
44
+
45
+
46
+@support_torch_compile
47
+class EagleMistralModel(MistralModel):
48
+ def __init__(
49
+ self,
50
+ *,
51
+ vllm_config: VllmConfig,
52
+ prefix: str = "",
53
+ start_layer_id: int = 0,
54
+ ) -> None:
55
+ # Bypass MistralModel.__init__ to avoid creating duplicate attention
56
+ # layer entries in the global context.
57
+ nn.Module.__init__(self)
58
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
59
+ self.vocab_size = self.config.vocab_size
60
+ # Get drafter's quantization config
61
+ self.quant_config = get_draft_quant_config(vllm_config)
62
+
63
+ self.embed_tokens = VocabParallelEmbedding(
64
+ self.config.vocab_size,
65
+ self.config.hidden_size,
66
+ prefix=maybe_prefix(prefix, "embed_tokens"),
67
+ quant_config=self.quant_config,
68
+ )
69
+
70
+ self.layers = nn.ModuleList(
71
+ [
72
+ EagleMistralDecoderLayer(
73
+ vllm_config,
74
+ prefix=maybe_prefix(prefix, f"layers.{i + start_layer_id}"),
75
+ config=self.config,
76
+ )
77
+ for i in range(self.config.num_hidden_layers)
78
+ ]
79
+ )
80
+ self.fc = RowParallelLinear(
81
+ self.config.hidden_size * 2,
82
+ self.config.hidden_size,
83
+ bias=False,
84
+ input_is_parallel=False,
85
+ quant_config=self.quant_config,
86
+ prefix=maybe_prefix(prefix, "fc"),
87
+ return_bias=False,
88
+ )
89
+ self.norm = RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps)
90
+
91
+ def forward(
92
+ self,
93
+ input_ids: torch.Tensor,
94
+ positions: torch.Tensor,
95
+ hidden_states: torch.Tensor,
96
+ inputs_embeds: torch.Tensor | None = None,
97
+ ) -> tuple[torch.Tensor, torch.Tensor]:
98
+ if inputs_embeds is None:
99
+ inputs_embeds = self.embed_input_ids(input_ids)
100
+ hidden_states = self.fc(torch.cat((inputs_embeds, hidden_states), dim=-1))
101
+ residual = None
102
+ for layer in self.layers:
103
+ hidden_states, residual = layer(
104
+ positions,
105
+ hidden_states,
106
+ residual,
107
+ )
108
+ hidden_states, _ = self.norm(hidden_states, residual)
109
+ return hidden_states, hidden_states
110
+
111
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
112
+ # Pretend embed_tokens is loaded; the actual weight is shared
113
+ # from the target model at runtime by `load_eagle_model`.
114
+ return super().load_weights(weights) | {"embed_tokens.weight"}
115
+
116
+
117
+class EagleMistralForCausalLM(MistralForCausalLM):
118
+ mistral_mapping = MistralForCausalLM.mistral_mapping | {
119
+ "eagle_linear": "model.fc",
120
+ }
121
+
122
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
123
+ # Bypass MistralForCausalLM.__init__ to use the draft model config
124
+ # and to avoid creating an lm_head.
125
+ nn.Module.__init__(self)
126
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
127
+ target_layer_num = vllm_config.model_config.get_num_layers(
128
+ vllm_config.parallel_config
129
+ )
130
+ self.model = EagleMistralModel(
131
+ vllm_config=vllm_config, prefix="model", start_layer_id=target_layer_num
132
+ )
133
+
134
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
135
+ self.logits_processor = LogitsProcessor(
136
+ self.config.vocab_size, scale=logit_scale
137
+ )
138
+
139
+ def forward(
140
+ self,
141
+ input_ids: torch.Tensor,
142
+ positions: torch.Tensor,
143
+ hidden_states: torch.Tensor,
144
+ inputs_embeds: torch.Tensor | None = None,
145
+ ) -> tuple[torch.Tensor, torch.Tensor]:
146
+ return self.model(input_ids, positions, hidden_states, inputs_embeds)
147
+
148
+ def embed_input_ids(
149
+ self,
150
+ input_ids: torch.Tensor,
151
+ multimodal_embeddings: MultiModalEmbeddings | None = None,
152
+ *,
153
+ is_multimodal: torch.Tensor | None = None,
154
+ ) -> torch.Tensor:
155
+ inputs_embeds = super().embed_input_ids(input_ids)
156
+
157
+ if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
158
+ return inputs_embeds
159
+
160
+ assert is_multimodal is not None
161
+
162
+ return _merge_multimodal_embeddings(
163
+ inputs_embeds=inputs_embeds,
164
+ multimodal_embeddings=multimodal_embeddings,
165
+ is_multimodal=is_multimodal,
166
+ )
vllm/model_executor/models/registry.py
+1
@@ -584,6 +584,7 @@ _SPECULATIVE_DECODING_MODELS = {
584
"LlamaForCausalLMEagle3": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
585
"Eagle3Qwen2_5vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
586
"Eagle3Qwen3vlForCausalLM": ("llama_eagle3", "Eagle3LlamaForCausalLM"),
587
+ "EagleMistralForCausalLM": ("mistral_eagle", "EagleMistralForCausalLM"),
588
"EagleMistralLarge3ForCausalLM": (
589
"mistral_large_3_eagle",
590
"EagleMistralLarge3ForCausalLM",