xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-05-30 09:17:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Speculators][Speculative Decoding] Fix Kimi K2 Eagle3 Support

Browse Source 
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
This commit is contained in:
chaunceyjiang 2025-12-24 03:05:37 +00:00
parent 27c6c2f98c
commit 3a4929e9ce
1 changed files with 20 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
vllm/model_executor/models/deepseek_v2.py
View File

@ -85,7 +85,13 @@ from vllm.v1.attention.backends.mla.indexer import (
from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec from vllm.v1.kv_cache_interface import KVCacheSpec, MLAAttentionSpec
from vllm.v1.worker.workspace import current_workspace_manager from vllm.v1.worker.workspace import current_workspace_manager
from .interfaces import MixtureOfExperts, SupportsEagle, SupportsLoRA, SupportsPP from .interfaces import (
MixtureOfExperts,
SupportsEagle,
SupportsEagle3,
SupportsLoRA,
SupportsPP,
)
from .utils import ( from .utils import (
PPMissingLayer, PPMissingLayer,
is_pp_missing_parameter, is_pp_missing_parameter,
@ -1378,7 +1384,12 @@ class DeepseekV2MixtureOfExperts(MixtureOfExperts):
class DeepseekV2ForCausalLM( class DeepseekV2ForCausalLM(
nn.Module, SupportsPP, DeepseekV2MixtureOfExperts, SupportsLoRA, SupportsEagle nn.Module,
SupportsPP,
DeepseekV2MixtureOfExperts,
SupportsLoRA,
SupportsEagle,
SupportsEagle3,
): ):
packed_modules_mapping = { packed_modules_mapping = {
"gate_up_proj": ["gate_proj", "up_proj"], "gate_up_proj": ["gate_proj", "up_proj"],
@ -1460,6 +1471,13 @@ class DeepseekV2ForCausalLM(
def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.embed_input_ids(input_ids) return self.model.embed_input_ids(input_ids)
def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
self.model.aux_hidden_state_layers = layers
def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
num_layers = len(self.model.layers)
return (2, num_layers // 2, num_layers - 3)
def forward( def forward(
self, self,
input_ids: torch.Tensor, input_ids: torch.Tensor,