diff --git a/python/sglang/srt/models/longcat_flash.py b/python/sglang/srt/models/longcat_flash.py index f3860e572..3530609ba 100644 --- a/python/sglang/srt/models/longcat_flash.py +++ b/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]