feat: longcat flash add aux layers capture for eagle3 (#14161)

python/sglang/srt/models/longcat_flash.py

@@ -32,7 +32,7 @@
 
 import concurrent.futures
 import logging
-from typing import Iterable, Optional, Tuple
+from typing import Iterable, List, Optional, Tuple
 
 import torch
 from torch import nn
@@ -511,6 +511,7 @@ class LongcatFlashModel(nn.Module):
             ]
         )
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.layers_to_capture = []
 
     def get_input_embeddings(self) -> torch.Tensor:
         return self.embed_tokens
@@ -536,7 +537,10 @@ class LongcatFlashModel(nn.Module):
 
         residual = None
 
+        aux_hidden_states = []
         for i in range(total_num_layers):
+            if i in self.layers_to_capture:
+                aux_hidden_states.append(hidden_states + residual)
             with get_global_expert_distribution_recorder().with_current_layer(i):
                 layer = self.layers[i]
                 hidden_states, residual = layer(
@@ -548,7 +552,11 @@ class LongcatFlashModel(nn.Module):
                 hidden_states = self.norm(hidden_states)
             else:
                 hidden_states, _ = self.norm(hidden_states, residual)
-        return hidden_states
+
+        if len(aux_hidden_states) == 0:
+            return hidden_states
+
+        return hidden_states, aux_hidden_states
 
 
 class LongcatFlashForCausalLM(nn.Module):
@@ -588,6 +596,7 @@ class LongcatFlashForCausalLM(nn.Module):
             use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
         )
         self.logits_processor = LogitsProcessor(config)
+        self.capture_aux_hidden_states = False
 
     def get_input_embeddings(self) -> nn.Embedding:
         return self.model.embed_tokens
@@ -602,8 +611,12 @@ class LongcatFlashForCausalLM(nn.Module):
     ) -> torch.Tensor:
         hidden_states = self.model(input_ids, positions, forward_batch, input_embeds)
 
+        aux_hidden_states = None
+        if self.capture_aux_hidden_states:
+            hidden_states, aux_hidden_states = hidden_states
+
         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
         )
 
     def post_load_weights(self, weight_names=None):
@@ -1023,5 +1036,14 @@ class LongcatFlashForCausalLM(nn.Module):
             num_logical_experts=config.n_routed_experts,
         )
 
+    def set_eagle3_layers_to_capture(self, layer_ids: Optional[List[int]] = None):
+        if layer_ids is None:
+            self.capture_aux_hidden_states = True
+            num_layers = self.config.num_hidden_layers
+            self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3]
+        else:
+            self.capture_aux_hidden_states = True
+            self.model.layers_to_capture = [val + 1 for val in layer_ids]
+
 
 EntryClass = [LongcatFlashForCausalLM]