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
Eagle3 that supports the Minicpm3 model (#24243)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: liudan <adan@minicpm.com> Co-authored-by: liudan <liudan@qq.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Lucia Fang <116399278+luccafong@users.noreply.github.com>
阿丹(adan) committed
Sep 27, 2025 at 01:04 UTC
33f6aaf9725ab3c0b07ed88d8a6113621f9eb4d2
2 files changed
+44
-9
vllm/config/speculative.py
+1
-1
@@ -540,7 +540,7 @@ class SpeculativeConfig:
540
"speculative decoding is > 1, but got "
541
f"{self.disable_by_batch_size=}")
542
543
- eagle3_target_supported = ["llama", "qwen", "gpt_oss"]
543
+ eagle3_target_supported = ["llama", "qwen", "minicpm", "gpt_oss"]
544
if self.method == "eagle3" and self.target_model_config and not any(
545
supported_model in
546
self.target_model_config.hf_text_config.model_type
vllm/model_executor/models/minicpm.py
+43
-8
@@ -55,7 +55,7 @@ from vllm.model_executor.utils import set_weight_attrs
55
from vllm.platforms import current_platform
56
from vllm.sequence import IntermediateTensors
57
58
-from .interfaces import SupportsLoRA, SupportsPP
58
+from .interfaces import SupportsEagle3, SupportsLoRA, SupportsPP
59
from .utils import (AutoWeightsLoader, is_pp_missing_parameter,
60
make_empty_intermediate_tensors_factory, make_layers,
61
maybe_prefix)
@@ -381,6 +381,9 @@ class MiniCPMModel(nn.Module):
381
self.num_experts = getattr(self.config, "num_experts", 0)
382
self._init_layers(prefix, config, cache_config, quant_config)
383
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
384
+
385
+ self.aux_hidden_state_layers = tuple[int, ...]()
386
+
387
self.make_empty_intermediate_tensors = (
388
make_empty_intermediate_tensors_factory(
389
["hidden_states", "residual"], self.config.hidden_size))
@@ -408,7 +411,8 @@ class MiniCPMModel(nn.Module):
411
positions: torch.Tensor,
412
intermediate_tensors: Optional[IntermediateTensors] = None,
413
inputs_embeds: Optional[torch.Tensor] = None,
411
- ) -> Union[torch.Tensor, IntermediateTensors]:
414
+ ) -> Union[torch.Tensor, IntermediateTensors, tuple[torch.Tensor,
415
+ list[torch.Tensor]]]:
416
if get_pp_group().is_first_rank:
417
if inputs_embeds is not None:
418
hidden_states = inputs_embeds
@@ -419,18 +423,29 @@ class MiniCPMModel(nn.Module):
423
hidden_states = intermediate_tensors["hidden_states"]
424
residual = intermediate_tensors["residual"]
425
422
- for layer in islice(self.layers, self.start_layer, self.end_layer):
426
+ aux_hidden_states = []
427
+ for idx, layer in enumerate(
428
+ islice(self.layers, self.start_layer, self.end_layer)):
429
+ if idx in self.aux_hidden_state_layers:
430
+ aux_hidden_states.append(
431
+ hidden_states +
432
+ residual if residual is not None else hidden_states)
433
hidden_states, residual = layer(
434
positions,
435
hidden_states,
436
residual,
437
)
438
+
439
if not get_pp_group().is_last_rank:
440
return IntermediateTensors({
441
"hidden_states": hidden_states,
442
"residual": residual
443
})
444
+
445
hidden_states = self.norm(hidden_states)
446
+
447
+ if len(aux_hidden_states) > 0:
448
+ return hidden_states, aux_hidden_states
449
return hidden_states
450
451
def load_weights(self, weights: Iterable[tuple[str,
@@ -502,7 +517,7 @@ class MiniCPMModel(nn.Module):
517
return loaded_params
518
519
505
-class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
520
+class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP, SupportsEagle3):
521
packed_modules_mapping = {
522
"qkv_proj": [
523
"q_proj",
@@ -568,16 +583,36 @@ class MiniCPMForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
583
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
584
return self.model.get_input_embeddings(input_ids)
585
586
+ def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
587
+ self.model.aux_hidden_state_layers = layers
588
+
589
+ def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
590
+ num_layers = len(self.model.layers)
591
+ return (2, num_layers // 2, num_layers - 3)
592
+
593
def forward(
594
self,
595
input_ids: torch.Tensor,
596
positions: torch.Tensor,
597
intermediate_tensors: Optional[IntermediateTensors] = None,
598
inputs_embeds: Optional[torch.Tensor] = None,
577
- ) -> Union[torch.Tensor, IntermediateTensors]:
578
- hidden_states = self.model(input_ids, positions, intermediate_tensors,
579
- inputs_embeds) / self.scale_width
580
- return hidden_states
599
+ ) -> Union[torch.Tensor, IntermediateTensors, tuple[torch.Tensor,
600
+ list[torch.Tensor]]]:
601
+ model_output = self.model(input_ids, positions, intermediate_tensors,
602
+ inputs_embeds)
603
+
604
+ if isinstance(model_output, tuple) and len(model_output) == 2:
605
+ # Aux hidden states are present.
606
+ hidden_states, aux_hidden_states = model_output
607
+ hidden_states = hidden_states / self.scale_width
608
+ return hidden_states, aux_hidden_states
609
+ else:
610
+ # Only hidden states or IntermediateTensors
611
+ if isinstance(model_output, IntermediateTensors):
612
+ return model_output
613
+ else:
614
+ hidden_states = model_output / self.scale_width
615
+ return hidden_states
616
617
def compute_logits(
618
self,