diff --git a/python/sglang/srt/models/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py index cd5731671..1fe767cdb 100644 --- a/python/sglang/srt/models/gpt_oss.py +++ b/python/sglang/srt/models/gpt_oss.py @@ -720,15 +720,25 @@ class GptOssModel(nn.Module): hidden_states = pp_proxy_tensors["hidden_states"] residual = pp_proxy_tensors["residual"] + # Capture hidden-state boundaries: boundary 0 is the embedding output, + # and boundary i + 1 is the output after transformer block i. aux_hidden_states = [] + if self.start_layer in self.layers_to_capture: + aux_hidden_states.append( + hidden_states + residual if residual is not None else 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 + residual) layer = self.layers[i] hidden_states, residual = layer( positions, hidden_states, forward_batch, residual ) + if i + 1 in self.layers_to_capture: + aux_hidden_states.append( + hidden_states + residual + if residual is not None + else hidden_states + ) if not self.pp_group.is_last_rank: return PPProxyTensors( { @@ -1319,15 +1329,18 @@ class GptOssForCausalLM(nn.Module): if not self.pp_group.is_last_rank: return + num_layers = self.config.num_hidden_layers 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 - # we plus 1 here because in sglang, for the ith layer, it takes the output - # of the (i-1)th layer as aux hidden state - self.model.layers_to_capture = [val + 1 for val in layer_ids] + # Preserve IDs that already include the final hidden-state + # boundary; otherwise retain the legacy output-layer conversion. + if layer_ids and max(layer_ids) == num_layers: + self.model.layers_to_capture = list(layer_ids) + else: + self.model.layers_to_capture = [val + 1 for val in layer_ids] def set_dflash_layers_to_capture(self, layer_ids: List[int]): if not self.pp_group.is_last_rank: diff --git a/python/sglang/srt/models/llama_eagle3.py b/python/sglang/srt/models/llama_eagle3.py index 294710d11..be8e8a7c3 100644 --- a/python/sglang/srt/models/llama_eagle3.py +++ b/python/sglang/srt/models/llama_eagle3.py @@ -158,6 +158,21 @@ class LlamaModel(nn.Module): bias=getattr(config, "bias", False), ) + eagle_config = getattr(config, "eagle_config", None) or {} + # Normalize the concatenated target-model features once before the FC + # projection. This differs from fc_norm, which normalizes each feature + # chunk independently before concatenation. + self.norm_before_fc = bool( + eagle_config.get("norm_before_fc", getattr(config, "norm_before_fc", False)) + ) + if self.norm_before_fc: + self.input_norm = RMSNorm( + self.hidden_size_in * self.num_aux_hidden_states, + eps=config.rms_norm_eps, + ) + else: + self.input_norm = None + # Per-aux RMSNorm before fc; enabled via `fc_norm` or legacy `use_aux_norm` flag. use_fc_norm = getattr(config, "fc_norm", None) or getattr( config, "use_aux_norm", False @@ -212,6 +227,8 @@ class LlamaModel(nn.Module): hidden_states = forward_batch.spec_info.hidden_states if hidden_states.shape[-1] != embeds.shape[-1]: + if self.input_norm is not None: + hidden_states = self.input_norm(hidden_states) if self.fc_norm is not None: chunks = hidden_states.chunk(self.num_aux_hidden_states, dim=-1) hidden_states = torch.cat(