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 6, 21:45 UTC
Kimi k2.5 MLA based eagle3 (#36361)
Signed-off-by: Izzy Putterman <iputterman@nvidia.com> Signed-off-by: Jhao-Ting Chen <jhaotingc@nvidia.com> Co-authored-by: Izzy Putterman <iputterman@nvidia.com>
Jhao-Ting Chen committed
Mar 11, 2026 at 08:36 UTC
557389473755bff50b6d00c03ca5c68e5c37c9a0
8 files changed
+499
-8
tests/models/registry.py
+12
@@ -1137,6 +1137,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1137
speculative_model="yuhuili/EAGLE-LLaMA3-Instruct-8B",
1138
tokenizer="meta-llama/Meta-Llama-3-8B-Instruct",
1139
),
1140
+ "Eagle3DeepseekV2ForCausalLM": _HfExamplesInfo(
1141
+ "moonshotai/Kimi-K2.5",
1142
+ trust_remote_code=True,
1143
+ speculative_model="AQ-MedAI/Kimi-K25-eagle3",
1144
+ tokenizer="moonshotai/Kimi-K2.5",
1145
+ ),
1146
+ "Eagle3DeepseekV3ForCausalLM": _HfExamplesInfo(
1147
+ "moonshotai/Kimi-K2.5",
1148
+ trust_remote_code=True,
1149
+ speculative_model="AQ-MedAI/Kimi-K25-eagle3",
1150
+ tokenizer="moonshotai/Kimi-K2.5",
1151
+ ),
1152
"Eagle3LlamaForCausalLM": _HfExamplesInfo(
1153
"meta-llama/Llama-3.1-8B-Instruct",
1154
trust_remote_code=True,
vllm/config/speculative.py
+4
@@ -779,6 +779,10 @@ class SpeculativeConfig:
779
"hunyuan_v1_dense",
780
"afmoe",
781
"nemotron_h",
782
+ "deepseek_v2",
783
+ "deepseek_v3",
784
+ "kimi_k2",
785
+ "kimi_k25",
786
]
787
if (
788
self.method in ("eagle3", "extract_hidden_states")
vllm/model_executor/models/deepseek_eagle3.py
new
+419
@@ -0,0 +1,419 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+"""Eagle3 speculative decoding model for DeepseekV2/V3 with MLP (no MoE)."""
5
+
6
+import copy
7
+from collections.abc import Iterable
8
+
9
+import torch
10
+import torch.nn as nn
11
+from transformers import DeepseekV2Config, DeepseekV3Config
12
+
13
+from vllm.compilation.decorators import support_torch_compile
14
+from vllm.config import VllmConfig, get_current_vllm_config
15
+from vllm.logger import init_logger
16
+from vllm.model_executor.layers.layernorm import RMSNorm
17
+from vllm.model_executor.layers.linear import ReplicatedLinear
18
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
+from vllm.model_executor.layers.vocab_parallel_embedding import (
20
+ ParallelLMHead,
21
+ VocabParallelEmbedding,
22
+)
23
+from vllm.model_executor.model_loader.weight_utils import (
24
+ default_weight_loader,
25
+ maybe_remap_kv_scale_name,
26
+)
27
+from vllm.model_executor.models.deepseek_v2 import (
28
+ DeepseekV2ForCausalLM,
29
+ DeepseekV2MLAAttention,
30
+ DeepseekV2MLP,
31
+)
32
+from vllm.multimodal.inputs import NestedTensors
33
+
34
+from .utils import (
35
+ AutoWeightsLoader,
36
+ get_draft_quant_config,
37
+ maybe_prefix,
38
+ process_eagle_weight,
39
+)
40
+
41
+logger = init_logger(__name__)
42
+
43
+
44
+class DeepseekV2Eagle3DecoderLayer(nn.Module):
45
+ """
46
+ Eagle3 decoder layer for Deepseek that:
47
+ 1. Always uses MLP (not MoE)
48
+ 2. First layer accepts concatenated embeds + hidden_states
49
+ """
50
+
51
+ def __init__(
52
+ self,
53
+ vllm_config: VllmConfig,
54
+ prefix: str,
55
+ config: DeepseekV2Config | DeepseekV3Config | None = None,
56
+ layer_idx: int = 0,
57
+ ) -> None:
58
+ super().__init__()
59
+
60
+ if config is None:
61
+ config = vllm_config.model_config.hf_config
62
+ cache_config = vllm_config.cache_config
63
+ quant_config = get_draft_quant_config(vllm_config)
64
+
65
+ self.hidden_size = config.hidden_size
66
+ rope_scaling = getattr(config, "rope_scaling", None)
67
+ max_position_embeddings = getattr(config, "max_position_embeddings", 8192)
68
+
69
+ self.layer_idx = layer_idx
70
+
71
+ # MLA attention parameters
72
+ qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
73
+ qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
74
+ v_head_dim = getattr(config, "v_head_dim", 0)
75
+ kv_lora_rank = getattr(config, "kv_lora_rank", 0)
76
+ config = copy.copy(config)
77
+ if rope_scaling:
78
+ rope_params = rope_scaling.copy()
79
+ rope_params["rope_type"] = "deepseek_yarn"
80
+ else:
81
+ rope_params = {"rope_type": "default"}
82
+ config.rope_parameters = rope_params
83
+ self.self_attn = DeepseekV2MLAAttention(
84
+ vllm_config=vllm_config,
85
+ config=config,
86
+ hidden_size=self.hidden_size,
87
+ num_heads=config.num_attention_heads,
88
+ qk_nope_head_dim=qk_nope_head_dim,
89
+ qk_rope_head_dim=qk_rope_head_dim,
90
+ v_head_dim=v_head_dim,
91
+ q_lora_rank=config.q_lora_rank if hasattr(config, "q_lora_rank") else None,
92
+ kv_lora_rank=kv_lora_rank,
93
+ max_position_embeddings=max_position_embeddings,
94
+ cache_config=cache_config,
95
+ quant_config=quant_config,
96
+ prefix=f"{prefix}.self_attn",
97
+ input_size=2 * self.hidden_size if layer_idx == 0 else self.hidden_size,
98
+ )
99
+
100
+ # Always use MLP (not MoE) for Eagle3
101
+ self.mlp = DeepseekV2MLP(
102
+ hidden_size=config.hidden_size,
103
+ intermediate_size=config.intermediate_size,
104
+ hidden_act=config.hidden_act,
105
+ quant_config=quant_config,
106
+ prefix=f"{prefix}.mlp",
107
+ )
108
+
109
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
110
+ self.post_attention_layernorm = RMSNorm(
111
+ config.hidden_size, eps=config.rms_norm_eps
112
+ )
113
+
114
+ self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
115
+
116
+ if getattr(config, "norm_before_residual", False):
117
+ self._residual_norm = self._norm_before_residual
118
+ else:
119
+ self._residual_norm = self._norm_after_residual
120
+
121
+ def _norm_before_residual(
122
+ self, hidden_states: torch.Tensor
123
+ ) -> tuple[torch.Tensor, torch.Tensor]:
124
+ hidden_states = self.hidden_norm(hidden_states)
125
+ residual = hidden_states
126
+ return hidden_states, residual
127
+
128
+ def _norm_after_residual(
129
+ self, hidden_states: torch.Tensor
130
+ ) -> tuple[torch.Tensor, torch.Tensor]:
131
+ residual = hidden_states
132
+ hidden_states = self.hidden_norm(hidden_states)
133
+ return hidden_states, residual
134
+
135
+ def forward(
136
+ self,
137
+ positions: torch.Tensor,
138
+ embeds: torch.Tensor,
139
+ hidden_states: torch.Tensor,
140
+ residual: torch.Tensor | None,
141
+ ) -> tuple[torch.Tensor, torch.Tensor]:
142
+ if self.layer_idx == 0:
143
+ # First layer: concatenate embeds with hidden_states
144
+ embeds = self.input_layernorm(embeds)
145
+ hidden_states, residual = self._residual_norm(hidden_states=hidden_states)
146
+ hidden_states = torch.cat([embeds, hidden_states], dim=-1)
147
+ else:
148
+ # Subsequent layers: process hidden_states and residuals only
149
+ hidden_states, residual = self.input_layernorm(hidden_states, residual)
150
+
151
+ # Self Attention
152
+ hidden_states = self.self_attn(
153
+ positions=positions,
154
+ hidden_states=hidden_states,
155
+ llama_4_scaling=None,
156
+ )
157
+
158
+ hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
159
+
160
+ # Fully Connected (MLP, not MoE)
161
+ hidden_states = self.mlp(hidden_states)
162
+
163
+ return hidden_states, residual
164
+
165
+
166
+@support_torch_compile
167
+class DeepseekV2Eagle3Model(nn.Module):
168
+ def __init__(
169
+ self,
170
+ *,
171
+ vllm_config: VllmConfig,
172
+ start_layer_id: int = 0,
173
+ prefix: str = "",
174
+ ) -> None:
175
+ super().__init__()
176
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
177
+ self.vocab_size = self.config.vocab_size
178
+
179
+ # Get drafter's quantization config
180
+ self.quant_config = get_draft_quant_config(vllm_config)
181
+
182
+ current_vllm_config = get_current_vllm_config()
183
+
184
+ self.embed_tokens = VocabParallelEmbedding(
185
+ self.config.vocab_size,
186
+ self.config.hidden_size,
187
+ prefix=maybe_prefix(prefix, "embed_tokens"),
188
+ )
189
+
190
+ self.layers = nn.ModuleList(
191
+ [
192
+ DeepseekV2Eagle3DecoderLayer(
193
+ current_vllm_config,
194
+ prefix=maybe_prefix(prefix, f"layers.{layer_idx + start_layer_id}"),
195
+ config=self.config,
196
+ layer_idx=layer_idx,
197
+ )
198
+ for layer_idx in range(self.config.num_hidden_layers)
199
+ ]
200
+ )
201
+
202
+ # fc layer for combining auxiliary hidden states (3x hidden size input)
203
+ if hasattr(self.config, "target_hidden_size"):
204
+ fc_input_size = self.config.target_hidden_size * 3
205
+ else:
206
+ fc_input_size = self.config.hidden_size * 3
207
+
208
+ self.fc = ReplicatedLinear(
209
+ input_size=fc_input_size,
210
+ output_size=self.config.hidden_size,
211
+ bias=False,
212
+ params_dtype=vllm_config.model_config.dtype,
213
+ quant_config=self.quant_config,
214
+ prefix=maybe_prefix(prefix, "fc"),
215
+ return_bias=False,
216
+ )
217
+
218
+ self.norm = RMSNorm(
219
+ self.config.hidden_size,
220
+ eps=self.config.rms_norm_eps,
221
+ )
222
+
223
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
224
+ return self.embed_tokens(input_ids)
225
+
226
+ def forward(
227
+ self,
228
+ input_ids: torch.Tensor,
229
+ positions: torch.Tensor,
230
+ hidden_states: torch.Tensor,
231
+ input_embeds: torch.Tensor | None = None,
232
+ ) -> tuple[torch.Tensor, torch.Tensor]:
233
+ if input_embeds is None:
234
+ input_embeds = self.embed_input_ids(input_ids)
235
+ assert hidden_states.shape[-1] == input_embeds.shape[-1]
236
+
237
+ residual = None
238
+ for layer in self.layers:
239
+ hidden_states, residual = layer(
240
+ positions=positions,
241
+ embeds=input_embeds,
242
+ hidden_states=hidden_states,
243
+ residual=residual,
244
+ )
245
+ hidden_states, hidden_prenorm = self.norm(hidden_states, residual)
246
+ return hidden_states, hidden_prenorm
247
+
248
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
249
+ stacked_params_mapping = [
250
+ # (param_name, shard_name, shard_id)
251
+ (".gate_up_proj", ".gate_proj", 0),
252
+ (".gate_up_proj", ".up_proj", 1),
253
+ (".fused_qkv_a_proj", ".q_a_proj", 0),
254
+ (".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),
255
+ ]
256
+ params_dict = dict(self.named_parameters())
257
+ loaded_params: set[str] = set()
258
+
259
+ for name, loaded_weight in weights:
260
+ if "midlayer." in name:
261
+ name = name.replace("midlayer.", "layers.0.")
262
+
263
+ # Handle kv cache quantization scales
264
+ if self.quant_config is not None and (
265
+ scale_name := self.quant_config.get_cache_scale(name)
266
+ ):
267
+ param = params_dict[scale_name]
268
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
269
+ loaded_weight = (
270
+ loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
271
+ )
272
+ weight_loader(param, loaded_weight)
273
+ loaded_params.add(scale_name)
274
+ continue
275
+
276
+ # Remapping the name FP8 kv-scale
277
+ if "scale" in name:
278
+ name = maybe_remap_kv_scale_name(name, params_dict)
279
+ if name is None:
280
+ continue
281
+
282
+ for param_name, weight_name, shard_id in stacked_params_mapping:
283
+ if weight_name not in name:
284
+ continue
285
+ name = name.replace(weight_name, param_name)
286
+ param = params_dict[name]
287
+ weight_loader = param.weight_loader
288
+ weight_loader(param, loaded_weight, shard_id)
289
+ break
290
+ else:
291
+ if name not in params_dict:
292
+ continue
293
+ param = params_dict[name]
294
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
295
+ weight_loader(param, loaded_weight)
296
+ loaded_params.add(name)
297
+
298
+ return loaded_params
299
+
300
+
301
+class Eagle3DeepseekV2ForCausalLM(DeepseekV2ForCausalLM):
302
+ """Eagle3 speculative decoding model for DeepseekV2/V3."""
303
+
304
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
305
+ nn.Module.__init__(self)
306
+ self.config = vllm_config.speculative_config.draft_model_config.hf_config
307
+
308
+ # Ensure draft_vocab_size is set
309
+ if getattr(self.config, "draft_vocab_size", None) is None:
310
+ base_vocab_size = getattr(self.config, "vocab_size", None)
311
+ self.config.draft_vocab_size = base_vocab_size
312
+
313
+ target_layer_num = vllm_config.model_config.get_num_layers(
314
+ vllm_config.parallel_config
315
+ )
316
+
317
+ # Store target layer count in draft config
318
+ self.config.target_layer_count = target_layer_num
319
+
320
+ self.model = DeepseekV2Eagle3Model(
321
+ vllm_config=vllm_config, prefix="model", start_layer_id=target_layer_num
322
+ )
323
+
324
+ logit_scale = getattr(self.config, "logit_scale", 1.0)
325
+ self.lm_head = ParallelLMHead(
326
+ self.config.draft_vocab_size,
327
+ self.config.hidden_size,
328
+ prefix=maybe_prefix(prefix, "lm_head"),
329
+ )
330
+ self.logits_processor = LogitsProcessor(
331
+ self.config.draft_vocab_size, scale=logit_scale
332
+ )
333
+ self.draft_id_to_target_id = nn.Parameter(
334
+ torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
335
+ requires_grad=False,
336
+ )
337
+
338
+ def embed_input_ids(
339
+ self,
340
+ input_ids: torch.Tensor,
341
+ multimodal_embeddings: NestedTensors | None = None,
342
+ is_multimodal: torch.Tensor | None = None,
343
+ ) -> torch.Tensor:
344
+ return self.model.embed_input_ids(input_ids)
345
+
346
+ def forward(
347
+ self,
348
+ input_ids: torch.Tensor,
349
+ positions: torch.Tensor,
350
+ hidden_states: torch.Tensor,
351
+ inputs_embeds: torch.Tensor | None = None,
352
+ ) -> tuple[torch.Tensor, torch.Tensor]:
353
+ return self.model(input_ids, positions, hidden_states, inputs_embeds)
354
+
355
+ def compute_logits(
356
+ self,
357
+ hidden_states: torch.Tensor,
358
+ ) -> torch.Tensor | None:
359
+ logits = self.logits_processor(self.lm_head, hidden_states)
360
+ if self.draft_id_to_target_id is None:
361
+ assert logits.shape[1] == self.config.vocab_size, (
362
+ "Expected logits to have shape "
363
+ f"(*, {self.config.vocab_size}), but got {logits.shape}"
364
+ )
365
+ return logits
366
+
367
+ base = torch.arange(self.config.draft_vocab_size, device=logits.device)
368
+ targets = base + self.draft_id_to_target_id
369
+ logits_new = logits.new_full(
370
+ (
371
+ logits.shape[0],
372
+ self.config.vocab_size,
373
+ ),
374
+ float("-inf"),
375
+ )
376
+ logits_new[:, targets] = logits
377
+ return logits_new
378
+
379
+ def combine_hidden_states(
380
+ self,
381
+ hidden_states: torch.Tensor,
382
+ ) -> torch.Tensor:
383
+ # Combine multiple auxiliary hidden states returned by Eagle3
384
+ return self.model.fc(hidden_states)
385
+
386
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
387
+ model_weights = {}
388
+ includes_draft_id_mapping = False
389
+ includes_embed_tokens = False
390
+
391
+ for name, loaded_weight in weights:
392
+ if "t2d" in name:
393
+ continue
394
+ if "d2t" in name:
395
+ name = name.replace("d2t", "draft_id_to_target_id")
396
+ includes_draft_id_mapping = True
397
+ elif "lm_head" not in name:
398
+ name = "model." + name
399
+ if "embed_tokens" in name:
400
+ includes_embed_tokens = True
401
+ model_weights[name] = loaded_weight
402
+ process_eagle_weight(self, name)
403
+
404
+ skip_substrs = []
405
+ if not includes_draft_id_mapping:
406
+ skip_substrs.append("draft_id_to_target_id")
407
+ if not includes_embed_tokens:
408
+ skip_substrs.append("embed_tokens")
409
+
410
+ loader = AutoWeightsLoader(
411
+ self,
412
+ skip_prefixes=None,
413
+ skip_substrs=skip_substrs,
414
+ )
415
+ loader.load_weights(model_weights.items())
416
+
417
+
418
+# Aliases for compatibility
419
+Eagle3DeepseekV3ForCausalLM = Eagle3DeepseekV2ForCausalLM
vllm/model_executor/models/deepseek_v2.py
+39
-6
@@ -82,7 +82,13 @@ from vllm.v1.attention.backends.mla.indexer import (
82
)
83
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
84
85
-from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP
85
+from .interfaces import (
86
+ MixtureOfExperts,
87
+ SupportsEagle,
88
+ SupportsEagle3,
89
+ SupportsLoRA,
90
+ SupportsPP,
91
+)
92
from .utils import (
93
PPMissingLayer,
94
is_pp_missing_parameter,
@@ -828,6 +834,7 @@ class DeepseekV2MLAAttention(nn.Module):
834
quant_config: QuantizationConfig | None = None,
835
prefix: str = "",
836
topk_indices_buffer: torch.Tensor | None = None,
837
+ input_size: int | None = None,
838
) -> None:
839
super().__init__()
840
self.hidden_size = hidden_size
@@ -847,16 +854,20 @@ class DeepseekV2MLAAttention(nn.Module):
854
self.scaling = self.qk_head_dim**-0.5
855
self.max_position_embeddings = max_position_embeddings
856
857
+ # Use input_size for projection input dimensions if provided,
858
+ # otherwise default to hidden_size (used in Eagle3 Deepseek with MLA)
859
+ proj_input_size = input_size if input_size is not None else self.hidden_size
860
+
861
if self.q_lora_rank is not None:
862
self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
852
- self.hidden_size,
863
+ proj_input_size,
864
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
865
quant_config=quant_config,
866
prefix=f"{prefix}.fused_qkv_a_proj",
867
)
868
else:
869
self.kv_a_proj_with_mqa = ReplicatedLinear(
859
- self.hidden_size,
870
+ proj_input_size,
871
self.kv_lora_rank + self.qk_rope_head_dim,
872
bias=False,
873
quant_config=quant_config,
@@ -874,7 +885,7 @@ class DeepseekV2MLAAttention(nn.Module):
885
)
886
else:
887
self.q_proj = ColumnParallelLinear(
877
- self.hidden_size,
888
+ proj_input_size,
889
self.num_heads * self.qk_head_dim,
890
bias=False,
891
quant_config=quant_config,
@@ -1170,6 +1181,8 @@ class DeepseekV2Model(nn.Module):
1181
["hidden_states", "residual"], config.hidden_size
1182
)
1183
1184
+ self.aux_hidden_state_layers = tuple[int, ...]()
1185
+
1186
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1187
return self.embed_tokens(input_ids)
1188
@@ -1205,7 +1218,13 @@ class DeepseekV2Model(nn.Module):
1218
else:
1219
llama_4_scaling = None
1220
1208
- for layer in islice(self.layers, self.start_layer, self.end_layer):
1221
+ aux_hidden_states = []
1222
+ for idx, layer in enumerate(
1223
+ islice(self.layers, self.start_layer, self.end_layer),
1224
+ start=self.start_layer,
1225
+ ):
1226
+ if idx in self.aux_hidden_state_layers:
1227
+ aux_hidden_states.append(hidden_states + residual)
1228
hidden_states, residual = layer(
1229
positions, hidden_states, residual, llama_4_scaling
1230
)
@@ -1216,6 +1235,8 @@ class DeepseekV2Model(nn.Module):
1235
)
1236
1237
hidden_states, _ = self.norm(hidden_states, residual)
1238
+ if len(aux_hidden_states) > 0:
1239
+ return hidden_states, aux_hidden_states
1240
return hidden_states
1241
1242
@@ -1261,7 +1282,12 @@ class DeepseekV2MixtureOfExperts(MixtureOfExperts):
1282
1283
1284
class DeepseekV2ForCausalLM(
1264
- nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA, SupportsEagle
1285
+ nn.Module,
1286
+ SupportsPP,
1287
+ DeepseekV2MixtureOfExperts,
1288
+ SupportsLoRA,
1289
+ SupportsEagle,
1290
+ SupportsEagle3,
1291
):
1292
packed_modules_mapping = {
1293
"gate_up_proj": ["gate_proj", "up_proj"],
@@ -1340,6 +1366,13 @@ class DeepseekV2ForCausalLM(
1366
1367
self.extract_moe_parameters(example_moe)
1368
1369
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
1370
+ self.model.aux_hidden_state_layers = layers
1371
+
1372
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
1373
+ num_layers = len(self.model.layers)
1374
+ return (2, num_layers // 2, num_layers - 3)
1375
+
1376
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
1377
return self.model.embed_input_ids(input_ids)
1378
vllm/model_executor/models/kimi_k25.py
+14
-1
@@ -28,6 +28,8 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
28
CompressedTensorsConfig,
29
)
30
from vllm.model_executor.models.interfaces import (
31
+ SupportsEagle,
32
+ SupportsEagle3,
33
SupportsMultiModal,
34
SupportsPP,
35
SupportsQuant,
@@ -311,7 +313,12 @@ class KimiK25MultiModalProcessor(BaseMultiModalProcessor[KimiK25ProcessingInfo])
313
dummy_inputs=KimiK25DummyInputsBuilder,
314
)
315
class KimiK25ForConditionalGeneration(
314
- nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant
316
+ nn.Module,
317
+ SupportsMultiModal,
318
+ SupportsPP,
319
+ SupportsQuant,
320
+ SupportsEagle,
321
+ SupportsEagle3,
322
):
323
"""Kimi-K2.5 model for conditional generation.
324
@@ -480,6 +487,12 @@ class KimiK25ForConditionalGeneration(
487
logits = self.language_model.compute_logits(hidden_states)
488
return logits
489
490
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
491
+ self.language_model.set_aux_hidden_state_layers(layers)
492
+
493
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
494
+ return self.language_model.get_eagle3_aux_hidden_state_layers()
495
+
496
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
497
loader = AutoWeightsLoader(self)
498
return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
vllm/model_executor/models/registry.py
+2
@@ -551,6 +551,8 @@ _SPECULATIVE_DECODING_MODELS = {
551
"mistral_large_3_eagle",
552
"EagleMistralLarge3ForCausalLM",
553
),
554
+ "Eagle3DeepseekV2ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
555
+ "Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
556
"EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
557
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
558
"ErnieMTPModel": ("ernie_mtp", "ErnieMTP"),
vllm/transformers_utils/config.py
+1
@@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
87
funaudiochat="FunAudioChatConfig",
88
hunyuan_vl="HunYuanVLConfig",
89
isaac="IsaacConfig",
90
+ kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3
91
kimi_linear="KimiLinearConfig",
92
kimi_vl="KimiVLConfig",
93
kimi_k25="KimiK25Config",
vllm/v1/spec_decode/eagle.py
+8
-1
@@ -20,6 +20,7 @@ from vllm.logger import init_logger
20
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
21
from vllm.model_executor.model_loader import get_model
22
from vllm.model_executor.models import supports_multimodal
23
+from vllm.model_executor.models.deepseek_eagle3 import Eagle3DeepseekV2ForCausalLM
24
from vllm.model_executor.models.interfaces import SupportsMultiModal
25
from vllm.model_executor.models.llama_eagle3 import Eagle3LlamaForCausalLM
26
from vllm.multimodal import MULTIMODAL_REGISTRY
@@ -403,7 +404,9 @@ class SpecDecodeBaseProposer:
404
batch_size = common_attn_metadata.batch_size()
405
406
if self.method == "eagle3":
406
- assert isinstance(self.model, Eagle3LlamaForCausalLM)
407
+ assert isinstance(
408
+ self.model, (Eagle3LlamaForCausalLM, Eagle3DeepseekV2ForCausalLM)
409
+ )
410
target_hidden_states = self.model.combine_hidden_states(
411
target_hidden_states
412
)
@@ -1278,6 +1281,10 @@ class SpecDecodeBaseProposer:
1281
self.model.config.image_token_index = (
1282
target_model.config.vision_config.image_token_id
1283
)
1284
+ elif self.get_model_name(target_model) == "KimiK25ForConditionalGeneration":
1285
+ self.model.config.image_token_index = (
1286
+ target_model.config.media_placeholder_token_id
1287
+ )
1288
else:
1289
self.model.config.image_token_index = (
1290
target_model.config.image_token_index