chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: Add Ling Flash v2.0 support for Eagle3 (#15119)

Browse Source 
Co-authored-by: chenyefei.cyf <chenyefei.cyf@U-9V5T77LW-2356.local>
Co-authored-by: GeLee-Q <865038696@qq.com>
Co-authored-by: Gao016 <yngao016@163.com>
Co-authored-by: Shenggui Li <somerlee.9@gmail.com>
Co-authored-by: Yineng Zhang <me@zhyncs.com>
This commit is contained in:
yefei12
2026-02-01 15:57:45 +08:00
committed by GitHub
parent 3ca29dffc7
commit 855dd0546c
1 changed files with 30 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
python/sglang/srt/models/bailing_moe.py
View File

@@ -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