From 855dd0546ce7ba460f23bed08c469e63955228e2 Mon Sep 17 00:00:00 2001 From: yefei12 Date: Sun, 1 Feb 2026 15:57:45 +0800 Subject: [PATCH] feat: Add Ling Flash v2.0 support for Eagle3 (#15119) Co-authored-by: chenyefei.cyf Co-authored-by: GeLee-Q <865038696@qq.com> Co-authored-by: Gao016 Co-authored-by: Shenggui Li Co-authored-by: Yineng Zhang --- python/sglang/srt/models/bailing_moe.py | 31 ++++++++++++++++++++++++- 1 file changed, 30 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/models/bailing_moe.py b/python/sglang/srt/models/bailing_moe.py index fd637492c..a04e3d647 100644 --- a/python/sglang/srt/models/bailing_moe.py +++ b/python/sglang/srt/models/bailing_moe.py @@ -742,6 +742,8 @@ class BailingMoEModel(nn.Module): else: self.norm = PPMissingLayer(return_tuple=True) + self.layers_to_capture = [] + def forward( self, input_ids: torch.Tensor, @@ -764,6 +766,10 @@ class BailingMoEModel(nn.Module): aux_hidden_states = [] for i in range(self.start_layer, self.end_layer): with get_global_expert_distribution_recorder().with_current_layer(i): + if i in self.layers_to_capture: + aux_hidden_states.append( + hidden_states if residual is None else hidden_states + residual + ) layer = self.layers[i] hidden_states, residual = layer( positions, @@ -789,7 +795,10 @@ class BailingMoEModel(nn.Module): hidden_states = self.norm(hidden_states) else: hidden_states, _ = self.norm(hidden_states, residual) + + if len(aux_hidden_states) == 0: return hidden_states + return hidden_states, aux_hidden_states class BailingMoEForCausalLM(nn.Module): @@ -826,6 +835,8 @@ class BailingMoEForCausalLM(nn.Module): ) self.logits_processor = LogitsProcessor(config) + self.capture_aux_hidden_states = False + @property def start_layer(self): return self.model.start_layer @@ -863,9 +874,14 @@ class BailingMoEForCausalLM(nn.Module): input_embeds, pp_proxy_tensors=pp_proxy_tensors, ) + + aux_hidden_states = None + if self.capture_aux_hidden_states: + hidden_states, aux_hidden_states = hidden_states + if self.pp_group.is_last_rank: return self.logits_processor( - input_ids, hidden_states, self.lm_head, forward_batch + input_ids, hidden_states, self.lm_head, forward_batch, aux_hidden_states ) else: return hidden_states @@ -1014,6 +1030,19 @@ class BailingMoEForCausalLM(nn.Module): num_groups=None if num_groups == 0 else num_groups, ) + def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None): + if not self.pp_group.is_last_rank: + return + + self.capture_aux_hidden_states = True + if layer_ids is None: + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] + else: + # Add +1 because in SGLang, for the i-th layer, the auxiliary hidden state + # corresponds to the output of layer (i - 1). + self.model.layers_to_capture = [val + 1 for val in layer_ids] + class BailingMoeForCausalLM(BailingMoEForCausalLM): pass