[Speculative Decoding] Fix GPT-OSS EAGLE3 hidden states (#32334)

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,14 +1329,17 @@ 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
+            # Preserve IDs that already include the final hidden-state
-            # of the (i-1)th layer as aux 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]):

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(